datamaestro 0.8.1__py3-none-any.whl → 1.7.0__py3-none-any.whl

datamaestro/download/__init__.py

@@ -1,9 +1,116 @@
-from datamaestro.definitions import AbstractDataset, DatasetAnnotation, DatasetWrapper
+"""Resource system for dataset download and processing pipelines.
+
+This module defines the Resource interface and its concrete subclasses
+(FileResource, FolderResource, ValueResource) for managing dataset
+download and preprocessing steps as a directed acyclic graph (DAG).
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import shutil
+import warnings
+from abc import ABC, abstractmethod
+from enum import Enum
+from pathlib import Path
+from typing import IO, Union
+
+from attrs import define
+
+from datamaestro.definitions import AbstractDataset, DatasetAnnotation
 from datamaestro.utils import deprecated
 
+logger = logging.getLogger(__name__)
+
+# Module-level deprecation tracking (emit each category only once)
+_deprecation_warned: set[str] = set()
+
+
+def _warn_once(category: str, message: str):
+    """Emit a deprecation warning only once per category."""
+    if category not in _deprecation_warned:
+        _deprecation_warned.add(category)
+        warnings.warn(message, DeprecationWarning, stacklevel=3)
+
+
+# --- State metadata file helpers ---
+
+
+class ResourceStateFile:
+    """Manages the .state.json metadata file for resource states.
+
+    Location: <dataset.datapath>/.state.json
+
+    Format:
+        {
+          "version": 1,
+          "resources": {
+            "RESOURCE_NAME": {"state": "none"|"partial"|"complete"},
+            ...
+          }
+        }
+    """
+
+    VERSION = 1
+
+    def __init__(self, datapath: Path):
+        self._path = datapath / ".state.json"
+
+    def read(self, resource_name: str) -> "ResourceState":
+        """Read the state for a resource. Returns NONE if not found."""
+        data = self._load()
+        entry = data.get("resources", {}).get(resource_name)
+        if entry is None:
+            return ResourceState.NONE
+        return ResourceState(entry["state"])
+
+    def write(self, resource_name: str, state: "ResourceState"):
+        """Write the state for a resource (atomic write)."""
+        data = self._load()
+        if "resources" not in data:
+            data["resources"] = {}
+        data["resources"][resource_name] = {"state": state.value}
+        self._save(data)
+
+    def _load(self) -> dict:
+        if self._path.is_file():
+            with self._path.open("r") as f:
+                return json.load(f)
+        return {"version": self.VERSION, "resources": {}}
+
+    def _save(self, data: dict):
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        tmp = self._path.with_suffix(".tmp")
+        with tmp.open("w") as f:
+            json.dump(data, f, indent=2)
+        tmp.rename(self._path)
+
+
+# --- ResourceState enum ---
+
+
+class ResourceState(str, Enum):
+    """State of a resource in the preparation pipeline."""
+
+    NONE = "none"
+    """Not started / no data on disk."""
+
+    PARTIAL = "partial"
+    """Started but incomplete (error during download)."""
+
+    COMPLETE = "complete"
+    """Fully available."""
+
+
+# --- Lazy initialization decorator (backward compat) ---
+
 
 def initialized(method):
-    """Ensure the object is initialized"""
+    """Ensure the object is initialized (calls postinit on first use).
+
+    Deprecated: new Resource subclasses should not rely on this pattern.
+    """
 
     def wrapper(self, *args, **kwargs):
         if not self._post:
@@ -14,42 +121,521 @@ def initialized(method):
     return wrapper
 
 
-class Download(DatasetAnnotation):
-    """
-    Base class for all download handlers
+# --- SetupOptions (backward compat) ---
+
+
+@define(kw_only=True)
+class SetupOptions:
+    pass
+
+
+# --- Resource base class ---
+
+
+class Resource(DatasetAnnotation, ABC):
+    """Base class for all dataset resources.
+
+    A resource represents a single step in a dataset preparation pipeline.
+    Resources form a DAG: each resource declares its dependencies, and
+    the orchestrator ensures they are processed in topological order.
+
+    Usage modes:
+
+    1. Class attribute (preferred)::
+
+        @dataset(url="...")
+        class MyDataset(Base):
+            DATA = filedownloader("data.csv", "http://...", transient=True)
+            PROCESSED = SomeProcessor.from_file(DATA)
+
+    2. Decorator on function (deprecated, backward compat)::
+
+        @filedownloader("data.csv", "http://...")
+        @dataset(Base)
+        def my_dataset(data): ...
+
+    Two-path system:
+
+    - ``transient_path``: where download/processing writes data
+    - ``path``: final location after successful completion
+
+    The framework moves data from ``transient_path`` → ``path`` and then
+    marks the resource as COMPLETE. Subclass ``download()`` implementations
+    should always write to ``transient_path``.
+
+    State is persisted in a metadata file at::
+
+        <dataset.datapath>/.downloads/.state.json
     """
 
-    def __init__(self, varname: str):
-        self.varname = varname
-        # Ensures that the object is initialized
+    def __init__(
+        self,
+        varname: str | None = None,
+        *,
+        transient: bool = False,
+    ):
+        """
+        Args:
+            varname: Explicit resource name. If None, auto-set from
+                class attribute name during binding. Required when
+                used as a decorator (backward compat mode).
+            transient: If True, this resource's data can be deleted
+                after all its dependents reach COMPLETE.
+        """
+        self.name: str | None = varname
+        self._name_explicit: bool = varname is not None
+        self.dataset: AbstractDataset | None = None
+        self.transient: bool = transient
+        self._dependencies: list[Resource] = []
+        self._dependents: list[Resource] = []
+
+        # Backward compat: lazy initialization support
         self._post = False
 
-    def annotate(self, dataset: AbstractDataset):
-        # Register has a resource download
-        if self.varname in dataset.resources:
-            raise AssertionError("Name %s already declared as a resource", self.varname)
+    # ---- Properties ----
+
+    @property
+    def can_recover(self) -> bool:
+        """Whether partial downloads can be resumed.
+
+        When True and state is PARTIAL, existing data at transient_path
+        is preserved on error, allowing the next download() call to
+        resume from where it left off.
+
+        When False and state is PARTIAL, data at transient_path is
+        deleted and state is reset to NONE.
+
+        Default: False. Subclasses override to enable recovery.
+        """
+        return False
+
+    @property
+    def dependencies(self) -> list[Resource]:
+        """Resources that must be COMPLETE before this one can process.
+
+        Populated from constructor arguments. Subclasses with factory
+        methods should pass dependency resources to ``__init__`` and
+        store them in ``_dependencies``.
+        """
+        return self._dependencies
+
+    @property
+    def dependents(self) -> list[Resource]:
+        """Resources that depend on this one (inverse of dependencies).
+
+        Computed by the dataset after all resources are bound.
+        Used for eager transient cleanup decisions.
+        """
+        return self._dependents
+
+    @property
+    def path(self) -> Path:
+        """Final storage path for this resource's data.
+
+        This is where data lives after successful completion.
+        Default: ``dataset.datapath / self.name``
 
-        dataset.resources[self.varname] = self
-        self.definition = dataset
+        Subclasses may override to customize (e.g., add file extension).
+        """
+        return self.dataset.datapath / self.name
+
+    @property
+    def transient_path(self) -> Path:
+        """Temporary path where download/processing writes data.
+
+        During download(), subclasses write to this path.
+        After successful download, the framework moves the data from
+        transient_path to path, then marks state as COMPLETE.
+
+        Default: ``dataset.datapath / ".downloads" / self.name``
+        """
+        return self.dataset.datapath / ".downloads" / self.name
+
+    @property
+    def state(self) -> ResourceState:
+        """Current state, read from the metadata file.
+
+        If no metadata entry exists, returns NONE.
+        """
+        if self.dataset is None:
+            return ResourceState.NONE
+        state_file = ResourceStateFile(self.dataset.datapath)
+        return state_file.read(self.name)
+
+    @state.setter
+    def state(self, value: ResourceState) -> None:
+        """Update state in the metadata file (atomic write)."""
+        state_file = ResourceStateFile(self.dataset.datapath)
+        state_file.write(self.name, value)
 
     @property
     def context(self):
-        return self.definition.context
+        """Application context (from dataset)."""
+        return self.dataset.context
+
+    # ---- Abstract methods ----
+
+    @abstractmethod
+    def download(self, force: bool = False) -> None:
+        """Execute this resource's download/processing step.
+
+        Contract:
+
+        - Called only when all dependencies are COMPLETE.
+        - Must write output to ``self.transient_path``.
+        - The framework handles moving transient_path → path
+          and setting state to COMPLETE after this returns.
+        - If force=True, re-execute even if already COMPLETE.
+
+        Note: State management (COMPLETE/PARTIAL/NONE transitions,
+        moving transient_path → path) is handled by the framework,
+        NOT by the download() implementation.
+
+        Raises:
+            Exception: On download/processing failure. The framework
+                will handle PARTIAL state based on can_recover.
+        """
+        ...
+
+    @abstractmethod
+    def prepare(self):
+        """Return the value for dataset construction.
+
+        Called after download() has completed (state is COMPLETE).
+        Return type depends on the resource subclass:
+
+        - FileResource → Path
+        - FolderResource → Path
+        - ValueResource → resource-specific
+
+        For backward compat with function-based datasets, this value
+        is passed as a keyword argument to the dataset function.
+        """
+        ...
+
+    # ---- Concrete methods ----
+
+    def cleanup(self) -> None:
+        """Remove this resource's data from disk.
+
+        Called automatically for transient resources after all
+        dependents reach COMPLETE (eager cleanup).
+
+        Default implementation:
+
+        - Deletes self.path (file or directory)
+        - Deletes self.transient_path if it exists
+        - Sets self.state = NONE
+
+        Subclasses may override for custom cleanup.
+        """
+        for p in (self.path, self.transient_path):
+            if p.exists():
+                if p.is_dir():
+                    shutil.rmtree(p)
+                else:
+                    p.unlink()
+        self.state = ResourceState.NONE
+
+    def has_files(self) -> bool:
+        """Whether this resource produces files on disk.
+
+        Returns False for reference-only resources (e.g., links
+        to other datasets, in-memory values).
+        Default: True.
+        """
+        return True
+
+    # Backward compat alias
+    def hasfiles(self) -> bool:
+        """Deprecated: use has_files() instead."""
+        _warn_once("hasfiles", "hasfiles() is deprecated, use has_files()")
+        return self.has_files()
 
     def postinit(self):
+        """Legacy lazy initialization hook.
+
+        Deprecated: new Resource subclasses should perform
+        initialization in ``__init__`` or ``bind()``.
+        """
         pass
 
-    def hasfiles(self):
-        return True
+    # ---- Binding ----
+
+    def bind(self, name: str, dataset: AbstractDataset) -> None:
+        """Bind this resource to a dataset.
+
+        Called by the dataset class machinery during initialization.
+        Sets self.name (if not explicitly set via varname) and
+        self.dataset. Registers the resource in dataset.resources
+        and dataset.ordered_resources.
+
+        For class-based datasets: called by ``@dataset`` when it
+        processes class attributes.
+        For decorator-based: called by ``annotate()`` (existing protocol).
+        """
+        if not self._name_explicit:
+            self.name = name
+
+        assert self.dataset is None, (
+            f"Resource {self.name} is already bound to a dataset"
+        )
+
+        if self.name in dataset.resources:
+            raise AssertionError(f"Name {self.name} already declared as a resource")
+
+        dataset.resources[self.name] = self
+        dataset.ordered_resources.append(self)
+        self.dataset = dataset
+
+    def annotate(self, dataset: AbstractDataset) -> None:
+        """Register with a dataset (DatasetAnnotation protocol).
+
+        Deprecated for new code. Calls bind() internally.
+        """
+        _warn_once(
+            "annotate",
+            "Using resources as decorators is deprecated. "
+            "Define them as class attributes instead.",
+        )
+        self.bind(self.name, dataset)
+
+    def contextualize(self):
+        """When using an annotation inline, uses the current
+        dataset wrapper object.
+
+        Deprecated: use class-attribute resource definitions instead.
+        """
+        wrapper = AbstractDataset.processing()
+        self.annotate(wrapper)
+
+    def setup(
+        self,
+        dataset: Union[AbstractDataset],
+        options: SetupOptions = None,
+    ):
+        """Direct way to setup the resource (no annotation).
+
+        Deprecated: use class-attribute resource definitions instead.
+        """
+        self(dataset)
+        return self.prepare()
+
+    # ---- Factory pattern ----
+
+    @classmethod
+    def apply(cls, *args, **kwargs) -> "Resource":
+        """Factory classmethod for creating resource instances.
+
+        Allows defining shorthand factory functions::
+
+            filedownloader = FileDownloader.apply
+
+        Default implementation: ``return cls(*args, **kwargs)``
+        Subclasses may override for custom argument handling.
+        """
+        return cls(*args, **kwargs)
+
+    # ---- Backward compat: definition property ----
+
+    @property
+    def definition(self) -> AbstractDataset | None:
+        """Deprecated: use ``dataset`` attribute instead."""
+        _warn_once(
+            "definition",
+            "Resource.definition is deprecated, use Resource.dataset",
+        )
+        return self.dataset
+
+    # Backward compat: varname property
+    @property
+    def varname(self) -> str | None:
+        """Deprecated: use ``name`` attribute instead."""
+        _warn_once(
+            "varname",
+            "Resource.varname is deprecated, use Resource.name",
+        )
+        return self.name
+
+    @varname.setter
+    def varname(self, value: str | None):
+        self.name = value
+
+
+# --- FileResource ---
+
+
+class FileResource(Resource):
+    """A resource that produces a single file on disk.
+
+    Subclasses implement ``_download()`` to produce the file at the
+    given destination (which is ``self.transient_path``).
+    """
+
+    def __init__(
+        self,
+        filename: str,
+        *,
+        varname: str | None = None,
+        transient: bool = False,
+    ):
+        """
+        Args:
+            filename: The filename (with extension) for the produced file.
+                Used to construct the storage path.
+            varname: Explicit resource name. If None, derived from
+                filename (extension stripped) or class attribute name.
+            transient: See Resource.
+        """
+        import re
+
+        effective_varname = varname or re.sub(r"\..*$", "", filename)
+        super().__init__(varname=effective_varname, transient=transient)
+        # Only mark name as explicit if user actually passed varname
+        self._name_explicit = varname is not None
+        self.filename = filename
+
+    @property
+    def path(self) -> Path:
+        """Final path to the produced file.
+
+        ``dataset.datapath / self.filename``
+        """
+        return self.dataset.datapath / self.filename
+
+    @property
+    def transient_path(self) -> Path:
+        """Temporary path for writing during download.
+
+        ``dataset.datapath / ".downloads" / self.filename``
+        """
+        return self.dataset.datapath / ".downloads" / self.filename
+
+    def prepare(self) -> Path:
+        """Returns self.path."""
+        return self.path
+
+    def stream(self) -> IO[bytes] | None:
+        """Return a readable byte stream of the file content.
+
+        Returns None if streaming is not supported for this resource.
+        Default: returns None. Subclasses may override.
+
+        This allows downstream resources to consume data without
+        needing the file to be fully materialized on disk first.
+        """
+        return None
+
+    def download(self, force: bool = False) -> None:
+        """Downloads the file.
+
+        Delegates to ``_download(self.transient_path)``.
+        """
+        self._download(self.transient_path)
+
+    @abstractmethod
+    def _download(self, destination: Path) -> None:
+        """Subclass hook: download/produce the file at destination.
+
+        Args:
+            destination: The path to write the file to
+                (``self.transient_path``).
+        """
+        ...
+
+
+# --- FolderResource ---
+
+
+class FolderResource(Resource):
+    """A resource that produces a directory on disk.
+
+    Subclasses implement ``_download()`` to populate the directory at
+    the given destination (which is ``self.transient_path``).
+    """
+
+    @property
+    def path(self) -> Path:
+        """Final path to the produced directory.
+
+        ``dataset.datapath / self.name``
+        """
+        return self.dataset.datapath / self.name
+
+    @property
+    def transient_path(self) -> Path:
+        """Temporary path for writing during download.
+
+        ``dataset.datapath / ".downloads" / self.name``
+        """
+        return self.dataset.datapath / ".downloads" / self.name
+
+    def prepare(self) -> Path:
+        """Returns self.path."""
+        return self.path
+
+    def download(self, force: bool = False) -> None:
+        """Downloads/extracts the directory content to transient_path."""
+        self._download(self.transient_path)
+
+    @abstractmethod
+    def _download(self, destination: Path) -> None:
+        """Subclass hook: populate the directory at destination.
+
+        Args:
+            destination: The path to write to (``self.transient_path``).
+        """
+        ...
+
+
+# --- ValueResource ---
+
+
+class ValueResource(Resource):
+    """A resource that produces an in-memory value (no files on disk).
+
+    Used for resources like HuggingFace dataset handles that don't
+    produce local files. The transient_path/path two-path system
+    is not used; state tracking is still via metadata file.
+    """
+
+    def has_files(self) -> bool:
+        return False
+
+    @abstractmethod
+    def prepare(self):
+        """Return the in-memory value."""
+        ...
+
+
+# --- Deprecated compatibility classes ---
+
+
+class Download(Resource):
+    """Deprecated: use Resource instead."""
+
+    def __init_subclass__(cls):
+        _warn_once(
+            f"Download-{cls.__name__}",
+            f"Download is deprecated ({cls}): use `Resource`",
+        )
+        return super().__init_subclass__()
+
+
+# --- reference resource ---
 
-    def download(self, force=False):
-        """Downloads the content"""
-        raise NotImplementedError()
 
+class reference(Resource):
+    """References another dataset instead of downloading."""
 
-class reference(Download):
-    def __init__(self, varname, reference):
-        super().__init__(varname)
+    def __init__(self, varname=None, reference=None):
+        """
+        Args:
+            varname: The name of the variable.
+            reference: Another dataset to reference.
+        """
+        super().__init__(varname=varname)
+        assert reference is not None, "Reference cannot be null"
         self.reference = reference
 
     def prepare(self):
@@ -61,7 +647,7 @@ class reference(Download):
     def download(self, force=False):
         self.reference.__datamaestro__.download(force)
 
-    def hasfiles(self):
+    def has_files(self):
         # We don't really have files
         return False