Datasets¶

class AbstractDataset(ABC):
    """
    Abstract base class for type-based serialization.

    Subclasses auto-register via __init_subclass__. Supports lazy plugin discovery
    via entry points when a type/format combination is not found.

    Examples:
        Create a custom dataset using class parameters:

        >>> class JsonDataset(AbstractDataset, format="json", types=dict, extensions="json"):
        ...     def serialize(self, obj: dict) -> bytes:
        ...         import json
        ...
        ...         return json.dumps(obj).encode("utf-8")
        ...
        ...     def deserialize(self, data: bytes) -> dict:
        ...         import json
        ...
        ...         return json.loads(data.decode("utf-8"))

        Or using class variables:

        >>> class YamlDataset(AbstractDataset):
        ...     format = "yaml"
        ...     types = (dict,)
        ...     extensions = ("yaml", "yml")
        ...
        ...     def serialize(self, obj: dict) -> bytes: ...
        ...     def deserialize(self, data: bytes) -> dict: ...

        The dataset is now registered and can be retrieved:

        >>> dataset_class = AbstractDataset.get(dict, "json")
        >>> dataset = dataset_class()
        >>> dataset.serialize({"key": "value"})
        b'{"key": "value"}'
    """

    # Class-level registry: (type, format) -> Dataset class
    _registry: ClassVar[dict[tuple[type, str], type[AbstractDataset]]] = {}

    # Reverse lookup: extension -> list of (type, format) for Driver hints
    _extension_hints: ClassVar[dict[str, list[tuple[type, str]]]] = {}

    # Track which entry point modules we've already tried to load
    _discovered: ClassVar[set[str]] = set()

    # Whether to auto-discover plugins on lookup failure
    _auto_discover: ClassVar[bool] = True

    # Defined by subclasses via __init_subclass__ or as class variables
    format: ClassVar[str]
    types: ClassVar[tuple[type, ...] | Any]
    extensions: ClassVar[tuple[str, ...] | str | None]

    def __init_subclass__(
        cls,
        *,
        format: str | None = None,
        types: type | tuple[type, ...] | None = None,
        extensions: str | tuple[str, ...] | None = None,
        **kwargs: Any,
    ) -> None:
        """
        Auto-register subclasses in the dataset registry.

        Parameters can be provided as class parameters or as class variables.
        Class parameters take precedence over class variables.

        Args:
            format: Format identifier (e.g., 'pickle', 'parquet'). If None, checks for class
                variable. If neither exists, this is an abstract intermediate class and won't be
                registered.
            types: Python type(s) this dataset can serialize. Can be a single
                type or a tuple of types.
            extensions: Optional file extension(s) for format inference (without dots). Can be a
                single string or a tuple of strings.
            kwargs: Additional keyword arguments.

        Examples:
            Using class parameters:

            >>> class JsonDataset(AbstractDataset, format="json", types=dict): ...

            Using class variables:

            >>> class CsvDataset(AbstractDataset):
            ...     format = "csv"
            ...     types = (dict,)
            ...     extensions = ("csv",)
            ...     ...
        """
        super().__init_subclass__(**kwargs)

        # Resolve format: parameter > class variable > None (abstract class)
        resolved_format = format
        if resolved_format is None:
            resolved_format = getattr(cls, "format", None)

        if resolved_format is None:
            return

        # Resolve types: parameter > class variable > empty
        types = getattr(cls, "types", ()) if types is None else types
        if not types:
            raise DatasetError(
                f"Dataset class {cls.__name__} must specify 'types' either as a class parameter "
                f"or class variable"
            )
        resolved_types = (types,) if isinstance(types, type) else types

        # Resolve extensions: parameter > class variable > empty
        extensions = getattr(cls, "extensions", ()) if extensions is None else extensions
        extensions = () if extensions is None else extensions
        resolved_extensions = (extensions,) if isinstance(extensions, str) else extensions

        cls.format = resolved_format
        cls.supported_types = resolved_types  # type: ignore
        cls.file_extensions = resolved_extensions  # type: ignore

        # Register for each supported type
        for t in resolved_types:
            AbstractDataset._registry[(t, resolved_format)] = cls

        # Register extension hints for Driver
        for ext in resolved_extensions:
            if ext not in AbstractDataset._extension_hints:
                AbstractDataset._extension_hints[ext] = []
            for t in resolved_types:
                AbstractDataset._extension_hints[ext].append((t, resolved_format))

    def __init__(self, *args, **kwargs) -> None:
        # NOTE: Subclasses can define their own __init__ with specific parameters. The base
        # implementation allows arbitrary args/kwargs for inheritance, but raises an error if any
        # are provided to prevent silent bugs from unrecognized parameters.
        if args or kwargs:
            name = self.__class__.__name__
            raise ValueError(
                f"{name} does not accept parameters in its constructor, but got args={args} "
                f"kwargs={kwargs}"
            )

    @abstractmethod
    def serialize(self, obj: Any) -> bytes:
        """
        Convert a Python object to bytes.

        Args:
            obj: Python object to serialize.

        Returns:
            Serialized bytes representation.
        """
        ...

    @abstractmethod
    def deserialize(self, data: bytes) -> Any:
        """
        Convert bytes back to a value.

        Args:
            data: Serialized bytes to deserialize.

        Returns:
            The deserialized value.
        """
        ...

    @classmethod
    def get(cls, type_: type, format: str) -> type[AbstractDataset]:
        """
        Look up and return the Dataset type for a type/format combination.

        Supports lazy plugin discovery - if the combination isn't found,
        attempts to load relevant entry points before failing.

        Args:
            type_: The Python type to serialize/deserialize.
            format: The format identifier (e.g., 'pickle', 'parquet').

        Returns:
            An instantiated Dataset that can handle the type/format.

        Raises:
            ValueError: If no dataset is registered for the type/format.

        Examples:
            >>> dataset_cls = AbstractDataset.get(str, "text")
            >>> dataset = dataset_cls()
            >>> dataset.serialize("hello")
            b'hello'
        """
        key = (type_, format)

        # Fast path: exact match in registry
        if key in cls._registry:
            return cls._registry[key]

        # Check for subclass matches (e.g., custom dict subclass → dict handler)
        for (reg_type, reg_format), dataset_cls in cls._registry.items():
            if reg_format == format:
                try:
                    if issubclass(type_, reg_type):
                        return dataset_cls
                except TypeError:  # pragma: no cover - defensive
                    continue

        # Slow path: try discovering plugins for this type
        if cls._auto_discover:  # pragma: no branch - defensive
            cls._discover_for_type(type_)

            if key in cls._registry:
                return cls._registry[key]  # pragma: no cover - post-discovery fast path

            for (
                reg_type,
                reg_format,
            ), dataset_cls in cls._registry.items():  # pragma: no cover - post-discovery
                if reg_format == format:
                    try:
                        if issubclass(type_, reg_type):
                            return dataset_cls
                    except TypeError:
                        continue

        module = type_.__module__.split(".")[0]
        available_formats = cls.get_formats_for_type(type_)

        if available_formats:
            raise ValueError(
                f"No dataset registered for {type_.__name__} with format '{format}'.\n"
                f"Available formats for this type: {', '.join(sorted(available_formats))}"
            )
        else:
            raise ValueError(
                f"No dataset registered for {type_.__name__} with format '{format}'.\n"
                f"Try: pip install daglite-datasets[{module}]"
            )

    @classmethod
    def infer_format(cls, type_: type, extension: str | None = None) -> str:
        """
        Infer the format for a type, optionally using extension as a hint.

        Args:
            type_: The Python type.
            extension: Optional file extension (without dot) to help infer format.

        Returns:
            The inferred format identifier.

        Raises:
            ValueError: If no format can be inferred for the type.

        Examples:
            >>> AbstractDataset.infer_format(str, "txt")
            'text'
            >>> AbstractDataset.infer_format(dict, "pkl")
            'pickle'
        """
        # Try extension hint first
        if extension:
            hints = cls._extension_hints.get(extension, [])
            for hint_type, hint_format in hints:
                if hint_type == type_:
                    return hint_format
            # Check subclasses
            for hint_type, hint_format in hints:
                try:
                    if issubclass(type_, hint_type):
                        return hint_format
                except TypeError:  # pragma: no cover - defensive
                    continue

        # Fall back to first registered format for this type
        for (reg_type, reg_format), _ in cls._registry.items():
            if reg_type == type_:
                return reg_format

        # Check subclass matches
        for (reg_type, reg_format), _ in cls._registry.items():
            try:
                if issubclass(type_, reg_type):
                    return reg_format
            except TypeError:  # pragma: no cover - defensive
                continue

        # Try auto-discovery
        if cls._auto_discover:  # pragma: no branch - defensive
            cls._discover_for_type(type_)

            for (reg_type, reg_format), _ in cls._registry.items():
                if reg_type == type_:
                    return reg_format  # pragma: no cover - post-discovery format match

        raise ValueError(f"No default format registered for {type_.__name__}")

    @classmethod
    def get_formats_for_type(cls, type_: type) -> set[str]:
        """
        Get all registered formats for a type.

        Args:
            type_: The Python type.

        Returns:
            Set of format identifiers registered for this type.

        Examples:
            >>> "text" in AbstractDataset.get_formats_for_type(str)
            True
        """
        formats: set[str] = set()
        for (reg_type, reg_format), _ in cls._registry.items():
            if reg_type == type_:
                formats.add(reg_format)
            else:
                try:
                    if issubclass(type_, reg_type):
                        formats.add(reg_format)
                except TypeError:  # pragma: no cover - defensive
                    continue
        return formats

    @classmethod
    def get_extension(cls, type_: type, format: str) -> str | None:
        """
        Get the preferred file extension for a type/format combination.

        Args:
            type_: The Python type.
            format: The format identifier.

        Returns:
            The preferred file extension (without dot), or None if not found.

        Examples:
            >>> AbstractDataset.get_extension(str, "text")
            'txt'
        """
        key = (type_, format)
        if key in cls._registry:
            dataset_cls = cls._registry[key]
            if dataset_cls.file_extensions:
                return dataset_cls.file_extensions[0]
        return None

    @classmethod
    def load_plugins(cls, *names: str, auto_discover: bool | None = None) -> None:
        """
        Explicitly load dataset plugins.

        Args:
            *names: Plugin names to load (e.g., "pandas", "numpy").
                If empty, loads all installed plugins.
            auto_discover: If provided, sets whether lazy discovery is enabled.
                Set to False to disable automatic plugin loading on lookup failure.

        Examples:
            Load specific plugins:

            >>> AbstractDataset.load_plugins("pandas", "numpy")

            Load all installed plugins:

            >>> AbstractDataset.load_plugins()

            Disable auto-discovery (strict mode):

            >>> AbstractDataset.load_plugins("pandas", auto_discover=False)
        """
        if auto_discover is not None:
            cls._auto_discover = auto_discover

        try:
            eps = importlib.metadata.entry_points(group="daglite.datasets")
        except TypeError:  # pragma: no cover
            # Python < 3.10 compatibility
            all_eps = importlib.metadata.entry_points()
            eps = all_eps.get("daglite.datasets", [])  # type: ignore[assignment]

        for ep in eps:
            if not names or ep.name in names:  # pragma: no cover - requires plugins
                try:
                    ep.load()  # Imports module, triggers __init_subclass__
                    cls._discovered.add(ep.name)
                except Exception:  # pragma: no cover
                    pass  # Plugin failed to load, continue with others

    @classmethod
    def _discover_for_type(cls, type_: type) -> None:
        """
        Lazily discover and load plugins that might handle this type.

        Uses the type's module name to find relevant entry points.
        For example, pandas.DataFrame → looks for "pandas" entry point.
        """
        module_name = type_.__module__.split(".")[0]

        if module_name in cls._discovered:
            return

        cls._discovered.add(module_name)

        try:
            eps = importlib.metadata.entry_points(group="daglite.datasets")
        except TypeError:  # pragma: no cover
            # Python < 3.10 compatibility
            all_eps = importlib.metadata.entry_points()
            eps = all_eps.get("daglite.datasets", [])  # type: ignore[assignment]

        for ep in eps:
            if ep.name == module_name:  # pragma: no cover - requires plugins
                try:
                    ep.load()
                except Exception:  # pragma: no cover
                    pass  # Plugin failed to load
                break

class DatasetStore:
    """
    High-level store that handles serialization via Datasets.

    This wraps a Driver (like FileDriver) and adds automatic
    serialization/deserialization using the Dataset registry.

    Format is inferred from the file extension in the key.

    This is the user-facing API - it accepts Python objects and handles
    all serialization internally.

    Examples:
        >>> from daglite.datasets import DatasetStore
        >>> store = DatasetStore("/tmp/outputs")  # doctest: +SKIP
        >>> store.save("data.pkl", {"data": [1, 2, 3]})  # doctest: +SKIP
        >>> store.load("data.pkl", dict)  # doctest: +SKIP
        {'data': [1, 2, 3]}
    """

    def __init__(self, driver: Driver | str) -> None:
        """
        Initialize with a driver.

        Args:
            driver: A Driver instance or string path (creates FileDriver).
        """
        self._driver: Driver
        if isinstance(driver, str):
            from daglite.drivers import FileDriver

            self._driver = FileDriver(driver)
        else:
            self._driver = driver

    @property
    def base_path(self) -> str:
        """Get base path (for FileDriver compatibility)."""
        return getattr(self._driver, "base_path", "")

    @property
    def is_local(self) -> bool:
        """Whether the underlying driver accesses local storage."""
        return getattr(self._driver, "is_local", True)

    def save(
        self,
        key: str,
        value: Any,
        format: str | None = None,
        options: dict[str, Any] | None = None,
    ) -> str:
        """
        Save a value using Dataset serialization.

        Args:
            key: Storage key/path. Format hint from driver (e.g., extension).
            value: Value to serialize and save.
            format: Serialization format. If None, inferred from type and/or driver hint.
            options: Additional options passed to the Dataset's save method.

        Returns:
            The actual path where data was stored.
        """
        key = _resolve_key(key)
        value_type = type(value)

        if format is None:
            hint = self._driver.get_format_hint(key)
            format = AbstractDataset.infer_format(value_type, hint)

        dataset_cls = AbstractDataset.get(value_type, format)
        options = options or {}
        dataset = dataset_cls(**options)
        data = dataset.serialize(value)
        return self._driver.save(key, data)

    def load(
        self,
        key: str,
        return_type: type[T] | None = None,
        format: str | None = None,
        options: dict[str, Any] | None = None,
    ) -> T:
        """
        Load a value using Dataset deserialization.

        Args:
            key: Storage key/path. Format hint from driver (e.g., extension).
            return_type: Expected return type. If None, uses pickle format.
            format: Serialization format. If None, inferred from type and/or driver hint.
            options: Additional options passed to the Dataset's load method.

        Returns:
            The deserialized value.

        Raises:
            KeyError: If key not found.
        """
        key = _resolve_key(key)
        data = self._driver.load(key)

        if return_type is None and format is None:
            return pickle.loads(data)  # No type hint or format - assume pickle

        hint = self._driver.get_format_hint(key)
        if return_type is None:  # pragma: no cover
            raise ValueError(
                "return_type must be provided when 'format' is specified; "
                "omit 'format' and 'return_type' together to use pickle-based loading."
            )
        if not format:
            format = AbstractDataset.infer_format(return_type, hint)
        dataset_cls = AbstractDataset.get(return_type, format)
        options = options or {}
        dataset = dataset_cls(**options)
        return cast(T, dataset.deserialize(data))

    def exists(self, key: str) -> bool:
        """Check if a key exists."""
        return self._driver.exists(_resolve_key(key))

    def delete(self, key: str) -> None:
        """Delete stored data."""
        self._driver.delete(_resolve_key(key))

    def list_keys(self) -> list[str]:
        """List all stored keys."""
        return self._driver.list_keys()

    def __getstate__(self) -> dict[str, Any]:
        """Serialize for pickling."""
        return {"driver": self._driver}

    def __setstate__(self, state: dict[str, Any]) -> None:
        """Reconstruct from pickled state."""
        self._driver = state["driver"]