_providers.py

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


@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"""