atdata 0.2.2b1__py3-none-any.whl → 0.3.0b1__py3-none-any.whl

atdata/lens.py

@@ -14,30 +14,28 @@ Key components:
 Lenses support the functional programming concept of composable, well-behaved
 transformations that satisfy lens laws (GetPut and PutGet).
 
-Example:
-    ::
-
-        >>> @packable
-        ... class FullData:
-        ...     name: str
-        ...     age: int
-        ...     embedding: NDArray
-        ...
-        >>> @packable
-        ... class NameOnly:
-        ...     name: str
-        ...
-        >>> @lens
-        ... def name_view(full: FullData) -> NameOnly:
-        ...     return NameOnly(name=full.name)
-        ...
-        >>> @name_view.putter
-        ... def name_view_put(view: NameOnly, source: FullData) -> FullData:
-        ...     return FullData(name=view.name, age=source.age,
-        ...                     embedding=source.embedding)
-        ...
-        >>> ds = Dataset[FullData]("data.tar")
-        >>> ds_names = ds.as_type(NameOnly)  # Uses registered lens
+Examples:
+    >>> @packable
+    ... class FullData:
+    ...     name: str
+    ...     age: int
+    ...     embedding: NDArray
+    ...
+    >>> @packable
+    ... class NameOnly:
+    ...     name: str
+    ...
+    >>> @lens
+    ... def name_view(full: FullData) -> NameOnly:
+    ...     return NameOnly(name=full.name)
+    ...
+    >>> @name_view.putter
+    ... def name_view_put(view: NameOnly, source: FullData) -> FullData:
+    ...     return FullData(name=view.name, age=source.age,
+    ...                     embedding=source.embedding)
+    ...
+    >>> ds = Dataset[FullData]("data.tar")
+    >>> ds_names = ds.as_type(NameOnly)  # Uses registered lens
 """
 
 ##
@@ -56,23 +54,24 @@ from typing import (
     Optional,
     Generic,
     #
-    TYPE_CHECKING
+    TYPE_CHECKING,
 )
 
 if TYPE_CHECKING:
     from .dataset import PackableSample
 
 from ._protocols import Packable
+from ._exceptions import LensNotFoundError
 
 
 ##
 # Typing helpers
 
-DatasetType: TypeAlias = Type['PackableSample']
+DatasetType: TypeAlias = Type["PackableSample"]
 LensSignature: TypeAlias = Tuple[DatasetType, DatasetType]
 
-S = TypeVar( 'S', bound = Packable )
-V = TypeVar( 'V', bound = Packable )
+S = TypeVar("S", bound=Packable)
+V = TypeVar("V", bound=Packable)
 type LensGetter[S, V] = Callable[[S], V]
 type LensPutter[S, V] = Callable[[V, S], S]
 
@@ -80,7 +79,8 @@ type LensPutter[S, V] = Callable[[V, S], S]
 ##
 # Shortcut decorators
 
-class Lens( Generic[S, V] ):
+
+class Lens(Generic[S, V]):
     """A bidirectional transformation between two sample types.
 
     A lens provides a way to view and update data of type ``S`` (source) as if
@@ -92,22 +92,22 @@ class Lens( Generic[S, V] ):
         S: The source type, must derive from ``PackableSample``.
         V: The view type, must derive from ``PackableSample``.
 
-    Example:
-        ::
-
-            >>> @lens
-            ... def name_lens(full: FullData) -> NameOnly:
-            ...     return NameOnly(name=full.name)
-            ...
-            >>> @name_lens.putter
-            ... def name_lens_put(view: NameOnly, source: FullData) -> FullData:
-            ...     return FullData(name=view.name, age=source.age)
+    Examples:
+        >>> @lens
+        ... def name_lens(full: FullData) -> NameOnly:
+        ...     return NameOnly(name=full.name)
+        ...
+        >>> @name_lens.putter
+        ... def name_lens_put(view: NameOnly, source: FullData) -> FullData:
+        ...     return FullData(name=view.name, age=source.age)
     """
-    # TODO The above has a line for "Parameters:" that should be "Type Parameters:"; this is a temporary fix for `quartodoc` auto-generation bugs.
 
-    def __init__( self, get: LensGetter[S, V],
-                put: Optional[LensPutter[S, V]] = None
-            ) -> None:
+    # Note: The docstring uses "Parameters:" for type parameters as a workaround
+    # for quartodoc not supporting "Type Parameters:" sections.
+
+    def __init__(
+        self, get: LensGetter[S, V], put: Optional[LensPutter[S, V]] = None
+    ) -> None:
         """Initialize a lens with a getter and optional putter function.
 
         Args:
@@ -126,8 +126,8 @@ class Lens( Generic[S, V] ):
 
         # Check argument validity
 
-        sig = inspect.signature( get )
-        input_types = list( sig.parameters.values() )
+        sig = inspect.signature(get)
+        input_types = list(sig.parameters.values())
         if len(input_types) != 1:
             raise ValueError(
                 f"Lens getter must have exactly one parameter, got {len(input_types)}: "
@@ -135,7 +135,7 @@ class Lens( Generic[S, V] ):
             )
 
         # Update function details for this object as returned by annotation
-        functools.update_wrapper( self, get )
+        functools.update_wrapper(self, get)
 
         self.source_type: Type[Packable] = input_types[0].annotation
         self.view_type: Type[Packable] = sig.return_annotation
@@ -146,14 +146,15 @@ class Lens( Generic[S, V] ):
         # Determine and store the putter
         if put is None:
             # Trivial putter does not update the source
-            def _trivial_put( v: V, s: S ) -> S:
+            def _trivial_put(v: V, s: S) -> S:
                 return s
+
             put = _trivial_put
         self._putter = put
-    
+
     #
 
-    def putter( self, put: LensPutter[S, V] ) -> LensPutter[S, V]:
+    def putter(self, put: LensPutter[S, V]) -> LensPutter[S, V]:
         """Decorator to register a putter function for this lens.
 
         Args:
@@ -163,20 +164,18 @@ class Lens( Generic[S, V] ):
         Returns:
             The putter function, allowing this to be used as a decorator.
 
-        Example:
-            ::
-
-                >>> @my_lens.putter
-                ... def my_lens_put(view: ViewType, source: SourceType) -> SourceType:
-                ...     return SourceType(...)
+        Examples:
+            >>> @my_lens.putter
+            ... def my_lens_put(view: ViewType, source: SourceType) -> SourceType:
+            ...     return SourceType(field=view.field, other=source.other)
         """
         ##
         self._putter = put
         return put
-    
+
     # Methods to actually execute transformations
 
-    def put( self, v: V, s: S ) -> S:
+    def put(self, v: V, s: S) -> S:
         """Update the source based on a modified view.
 
         Args:
@@ -186,9 +185,9 @@ class Lens( Generic[S, V] ):
         Returns:
             An updated source of type ``S`` that reflects changes from the view.
         """
-        return self._putter( v, s )
+        return self._putter(v, s)
 
-    def get( self, s: S ) -> V:
+    def get(self, s: S) -> V:
         """Transform the source into the view type.
 
         Args:
@@ -197,14 +196,14 @@ class Lens( Generic[S, V] ):
         Returns:
             A view of the source as type ``V``.
         """
-        return self( s )
+        return self(s)
 
-    def __call__( self, s: S ) -> V:
+    def __call__(self, s: S) -> V:
         """Apply the lens transformation (same as ``get()``)."""
-        return self._getter( s )
+        return self._getter(s)
 
 
-def lens(  f: LensGetter[S, V] ) -> Lens[S, V]:
+def lens(f: LensGetter[S, V]) -> Lens[S, V]:
     """Decorator to create and register a lens transformation.
 
     This decorator converts a getter function into a ``Lens`` object and
@@ -218,19 +217,17 @@ def lens(  f: LensGetter[S, V] ) -> Lens[S, V]:
         A ``Lens[S, V]`` object that can be called to apply the transformation
         or decorated with ``@lens_name.putter`` to add a putter function.
 
-    Example:
-        ::
-
-            >>> @lens
-            ... def extract_name(full: FullData) -> NameOnly:
-            ...     return NameOnly(name=full.name)
-            ...
-            >>> @extract_name.putter
-            ... def extract_name_put(view: NameOnly, source: FullData) -> FullData:
-            ...     return FullData(name=view.name, age=source.age)
+    Examples:
+        >>> @lens
+        ... def extract_name(full: FullData) -> NameOnly:
+        ...     return NameOnly(name=full.name)
+        ...
+        >>> @extract_name.putter
+        ... def extract_name_put(view: NameOnly, source: FullData) -> FullData:
+        ...     return FullData(name=view.name, age=source.age)
     """
-    ret = Lens[S, V]( f )
-    _network.register( ret )
+    ret = Lens[S, V](f)
+    _network.register(ret)
     return ret
 
 
@@ -259,11 +256,11 @@ class LensNetwork:
 
     def __init__(self):
         """Initialize the lens registry (only on first instantiation)."""
-        if not hasattr(self, '_initialized'):  # Check if already initialized
+        if not hasattr(self, "_initialized"):  # Check if already initialized
             self._registry: Dict[LensSignature, Lens] = dict()
             self._initialized = True
-    
-    def register( self, _lens: Lens ):
+
+    def register(self, _lens: Lens):
         """Register a lens as the canonical transformation between two types.
 
         Args:
@@ -275,8 +272,8 @@ class LensNetwork:
             overwritten.
         """
         self._registry[_lens.source_type, _lens.view_type] = _lens
-    
-    def transform( self, source: DatasetType, view: DatasetType ) -> Lens:
+
+    def transform(self, source: DatasetType, view: DatasetType) -> Lens:
         """Look up the lens transformation between two sample types.
 
         Args:
@@ -293,12 +290,17 @@ class LensNetwork:
             Currently only supports direct transformations. Compositional
             transformations (chaining multiple lenses) are not yet implemented.
         """
-        ret = self._registry.get( (source, view), None )
+        ret = self._registry.get((source, view), None)
         if ret is None:
-            raise ValueError( f'No registered lens from source {source} to view {view}' )
+            available_targets = [
+                (sig[1], lens_obj.__name__)
+                for sig, lens_obj in self._registry.items()
+                if sig[0] is source and hasattr(lens_obj, "__name__")
+            ]
+            raise LensNotFoundError(source, view, available_targets)
 
         return ret
 
 
 # Global singleton registry instance
-_network = LensNetwork()
+_network = LensNetwork()

atdata/local/__init__.py

@@ -0,0 +1,71 @@
+"""Local storage backend for atdata datasets.
+
+Key classes:
+
+- ``Index``: Unified index with pluggable providers (SQLite default),
+  named repositories, and optional atmosphere backend.
+- ``LocalDatasetEntry``: Index entry with ATProto-compatible CIDs.
+- ``S3DataStore``: S3-compatible shard storage.
+"""
+
+from atdata.local._entry import (
+    LocalDatasetEntry,
+    BasicIndexEntry,
+    REDIS_KEY_DATASET_ENTRY,
+    REDIS_KEY_SCHEMA,
+)
+from atdata.local._schema import (
+    SchemaNamespace,
+    SchemaFieldType,
+    SchemaField,
+    LocalSchemaRecord,
+    _ATDATA_URI_PREFIX,
+    _LEGACY_URI_PREFIX,
+    _kind_str_for_sample_type,
+    _schema_ref_from_type,
+    _make_schema_ref,
+    _parse_schema_ref,
+    _increment_patch,
+    _python_type_to_field_type,
+    _build_schema_record,
+)
+from atdata.local._index import Index
+from atdata.local._s3 import (
+    S3DataStore,
+    _s3_env,
+    _s3_from_credentials,
+    _create_s3_write_callbacks,
+)
+from atdata.local._repo_legacy import Repo
+
+# Re-export third-party types that were previously importable from the
+# monolithic local.py (tests reference atdata.local.S3FileSystem, etc.)
+from s3fs import S3FileSystem  # noqa: F401 — re-exported for backward compat
+
+__all__ = [
+    # Public API
+    "Index",
+    "LocalDatasetEntry",
+    "BasicIndexEntry",
+    "S3DataStore",
+    "Repo",
+    "SchemaNamespace",
+    "SchemaFieldType",
+    "SchemaField",
+    "LocalSchemaRecord",
+    "REDIS_KEY_DATASET_ENTRY",
+    "REDIS_KEY_SCHEMA",
+    # Internal helpers (re-exported for backward compatibility)
+    "_ATDATA_URI_PREFIX",
+    "_LEGACY_URI_PREFIX",
+    "_kind_str_for_sample_type",
+    "_schema_ref_from_type",
+    "_make_schema_ref",
+    "_parse_schema_ref",
+    "_increment_patch",
+    "_python_type_to_field_type",
+    "_build_schema_record",
+    "_s3_env",
+    "_s3_from_credentials",
+    "_create_s3_write_callbacks",
+]

atdata/local/_entry.py

@@ -0,0 +1,157 @@
+"""Dataset entry model and Redis key constants."""
+
+from atdata._cid import generate_cid
+
+from dataclasses import dataclass, field
+from typing import Any, cast
+
+import msgpack
+from redis import Redis
+
+
+# Redis key prefixes for index entries and schemas
+REDIS_KEY_DATASET_ENTRY = "LocalDatasetEntry"
+REDIS_KEY_SCHEMA = "LocalSchema"
+
+
+@dataclass
+class LocalDatasetEntry:
+    """Index entry for a dataset stored in the local repository.
+
+    Implements the IndexEntry protocol for compatibility with AbstractIndex.
+    Uses dual identity: a content-addressable CID (ATProto-compatible) and
+    a human-readable name.
+
+    The CID is generated from the entry's content (schema_ref + data_urls),
+    ensuring the same data produces the same CID whether stored locally or
+    in the atmosphere. This enables seamless promotion from local to ATProto.
+
+    Attributes:
+        name: Human-readable name for this dataset.
+        schema_ref: Reference to the schema for this dataset.
+        data_urls: WebDataset URLs for the data.
+        metadata: Arbitrary metadata dictionary, or None if not set.
+    """
+
+    ##
+
+    name: str
+    """Human-readable name for this dataset."""
+
+    schema_ref: str
+    """Reference to the schema for this dataset."""
+
+    data_urls: list[str]
+    """WebDataset URLs for the data."""
+
+    metadata: dict | None = None
+    """Arbitrary metadata dictionary, or None if not set."""
+
+    _cid: str | None = field(default=None, repr=False)
+    """Content identifier (ATProto-compatible CID). Generated from content if not provided."""
+
+    # Legacy field for backwards compatibility during migration
+    _legacy_uuid: str | None = field(default=None, repr=False)
+    """Legacy UUID for backwards compatibility with existing Redis entries."""
+
+    def __post_init__(self):
+        """Generate CID from content if not provided."""
+        if self._cid is None:
+            self._cid = self._generate_cid()
+
+    def _generate_cid(self) -> str:
+        """Generate ATProto-compatible CID from entry content."""
+        # CID is based on schema_ref and data_urls - the identity of the dataset
+        content = {
+            "schema_ref": self.schema_ref,
+            "data_urls": self.data_urls,
+        }
+        return generate_cid(content)
+
+    @property
+    def cid(self) -> str:
+        """Content identifier (ATProto-compatible CID)."""
+        if self._cid is None:
+            raise RuntimeError(
+                "CID not initialized; this should not happen after __post_init__"
+            )
+        return self._cid
+
+    # Legacy compatibility
+
+    @property
+    def wds_url(self) -> str:
+        """Legacy property: returns first data URL for backwards compatibility."""
+        return self.data_urls[0] if self.data_urls else ""
+
+    @property
+    def sample_kind(self) -> str:
+        """Legacy property: returns schema_ref for backwards compatibility."""
+        return self.schema_ref
+
+    def write_to(self, redis: Redis):
+        """Persist this index entry to Redis.
+
+        Stores the entry as a Redis hash with key '{REDIS_KEY_DATASET_ENTRY}:{cid}'.
+
+        Args:
+            redis: Redis connection to write to.
+        """
+        save_key = f"{REDIS_KEY_DATASET_ENTRY}:{self.cid}"
+        data: dict[str, Any] = {
+            "name": self.name,
+            "schema_ref": self.schema_ref,
+            "data_urls": msgpack.packb(self.data_urls),  # Serialize list
+            "cid": self.cid,
+        }
+        if self.metadata is not None:
+            data["metadata"] = msgpack.packb(self.metadata)
+        if self._legacy_uuid is not None:
+            data["legacy_uuid"] = self._legacy_uuid
+
+        redis.hset(save_key, mapping=data)  # type: ignore[arg-type]
+
+    @classmethod
+    def from_redis(cls, redis: Redis, cid: str) -> "LocalDatasetEntry":
+        """Load an entry from Redis by CID.
+
+        Args:
+            redis: Redis connection to read from.
+            cid: Content identifier of the entry to load.
+
+        Returns:
+            LocalDatasetEntry loaded from Redis.
+
+        Raises:
+            KeyError: If entry not found.
+        """
+        save_key = f"{REDIS_KEY_DATASET_ENTRY}:{cid}"
+        raw_data = redis.hgetall(save_key)
+        if not raw_data:
+            raise KeyError(f"{REDIS_KEY_DATASET_ENTRY} not found: {cid}")
+
+        # Decode string fields, keep binary fields as bytes for msgpack
+        raw_data_typed = cast(dict[bytes, bytes], raw_data)
+        name = raw_data_typed[b"name"].decode("utf-8")
+        schema_ref = raw_data_typed[b"schema_ref"].decode("utf-8")
+        cid_value = raw_data_typed.get(b"cid", b"").decode("utf-8") or None
+        legacy_uuid = raw_data_typed.get(b"legacy_uuid", b"").decode("utf-8") or None
+
+        # Deserialize msgpack fields (stored as raw bytes)
+        data_urls = msgpack.unpackb(raw_data_typed[b"data_urls"])
+        metadata = None
+        if b"metadata" in raw_data_typed:
+            metadata = msgpack.unpackb(raw_data_typed[b"metadata"])
+
+        return cls(
+            name=name,
+            schema_ref=schema_ref,
+            data_urls=data_urls,
+            metadata=metadata,
+            _cid=cid_value,
+            _legacy_uuid=legacy_uuid,
+        )
+
+
+# Backwards compatibility alias
+BasicIndexEntry = LocalDatasetEntry