import os
import uuid
from abc import abstractmethod
from collections.abc import Callable
from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Protocol
[docs]
@dataclass
class PathInfo:
"""
Information about where and how to write a file.
:param directory_path: Directory into which files should be written
:param filename: Base filename to use generated by FilenameProvider, w/o extension
:param create_dir_depth: Optional depth of directories to create if they do not
exist
"""
directory_path: Path
filename: str
create_dir_depth: int = 0
class FilenameProvider(Protocol):
@abstractmethod
def __call__(self, device_name: str | None = None) -> str:
"""Get a filename to use for output data, w/o extension"""
class PathProvider(Protocol):
@abstractmethod
def __call__(self, device_name: str | None = None) -> PathInfo:
"""Get the current directory to write files into"""
class StaticFilenameProvider(FilenameProvider):
def __init__(self, filename: str):
self._static_filename = filename
def __call__(self, device_name: str | None = None) -> str:
return self._static_filename
class UUIDFilenameProvider(FilenameProvider):
def __init__(
self,
uuid_call_func: Callable = uuid.uuid4,
uuid_call_args: list | None = None,
):
self._uuid_call_func = uuid_call_func
self._uuid_call_args = uuid_call_args or []
def __call__(self, device_name: str | None = None) -> str:
if (
self._uuid_call_func in [uuid.uuid3, uuid.uuid5]
and len(self._uuid_call_args) < 2
):
raise ValueError(
f"To use {self._uuid_call_func} to generate UUID filenames,"
" UUID namespace and name must be passed as args!"
)
uuid_str = self._uuid_call_func(*self._uuid_call_args)
return f"{uuid_str}"
class AutoIncrementFilenameProvider(FilenameProvider):
def __init__(
self,
base_filename: str = "",
max_digits: int = 5,
starting_value: int = 0,
increment: int = 1,
inc_delimeter: str = "_",
):
self._base_filename = base_filename
self._max_digits = max_digits
self._current_value = starting_value
self._increment = increment
self._inc_delimeter = inc_delimeter
def __call__(self, device_name: str | None = None) -> str:
if len(str(self._current_value)) > self._max_digits:
raise ValueError(
f"Auto incrementing filename counter \
exceeded maximum of {self._max_digits} digits!"
)
padded_counter = f"{self._current_value:0{self._max_digits}}"
filename = f"{self._base_filename}{self._inc_delimeter}{padded_counter}"
self._current_value += self._increment
return filename
class StaticPathProvider(PathProvider):
def __init__(
self,
filename_provider: FilenameProvider,
directory_path: Path,
create_dir_depth: int = 0,
) -> None:
self._filename_provider = filename_provider
self._directory_path = directory_path
self._create_dir_depth = create_dir_depth
def __call__(self, device_name: str | None = None) -> PathInfo:
filename = self._filename_provider(device_name)
return PathInfo(
directory_path=self._directory_path,
filename=filename,
create_dir_depth=self._create_dir_depth,
)
class AutoIncrementingPathProvider(PathProvider):
def __init__(
self,
filename_provider: FilenameProvider,
base_directory_path: Path,
create_dir_depth: int = 0,
max_digits: int = 5,
starting_value: int = 0,
num_calls_per_inc: int = 1,
increment: int = 1,
inc_delimeter: str = "_",
base_name: str | None = None,
) -> None:
self._filename_provider = filename_provider
self._base_directory_path = base_directory_path
self._create_dir_depth = create_dir_depth
self._base_name = base_name
self._starting_value = starting_value
self._current_value = starting_value
self._num_calls_per_inc = num_calls_per_inc
self._inc_counter = 0
self._max_digits = max_digits
self._increment = increment
self._inc_delimeter = inc_delimeter
def __call__(self, device_name: str | None = None) -> PathInfo:
filename = self._filename_provider(device_name)
padded_counter = f"{self._current_value:0{self._max_digits}}"
auto_inc_dir_name = str(padded_counter)
if self._base_name is not None:
auto_inc_dir_name = (
f"{self._base_name}{self._inc_delimeter}{padded_counter}"
)
elif device_name is not None:
auto_inc_dir_name = f"{device_name}{self._inc_delimeter}{padded_counter}"
self._inc_counter += 1
if self._inc_counter == self._num_calls_per_inc:
self._inc_counter = 0
self._current_value += self._increment
return PathInfo(
directory_path=self._base_directory_path / auto_inc_dir_name,
filename=filename,
create_dir_depth=self._create_dir_depth,
)
class YMDPathProvider(PathProvider):
def __init__(
self,
filename_provider: FilenameProvider,
base_directory_path: Path,
create_dir_depth: int = -3, # Default to -3 to create YMD dirs
device_name_as_base_dir: bool = False,
) -> None:
self._filename_provider = filename_provider
self._base_directory_path = Path(base_directory_path)
self._create_dir_depth = create_dir_depth
self._device_name_as_base_dir = device_name_as_base_dir
def __call__(self, device_name: str | None = None) -> PathInfo:
sep = os.path.sep
current_date = date.today().strftime(f"%Y{sep}%m{sep}%d")
if device_name is None:
ymd_dir_path = current_date
elif self._device_name_as_base_dir:
ymd_dir_path = os.path.join(
current_date,
device_name,
)
else:
ymd_dir_path = os.path.join(
device_name,
current_date,
)
filename = self._filename_provider(device_name)
return PathInfo(
directory_path=self._base_directory_path / ymd_dir_path,
filename=filename,
create_dir_depth=self._create_dir_depth,
)
class NameProvider(Protocol):
@abstractmethod
def __call__(self) -> str:
"""Get the name to be used as a data_key in the descriptor document"""
class DatasetDescriber(Protocol):
@abstractmethod
async def np_datatype(self) -> str:
"""Represents the numpy datatype"""
@abstractmethod
async def shape(self) -> tuple[int, ...]:
"""Get the shape of the data collection"""
