atdata 0.2.3b1__py3-none-any.whl → 0.3.1b1__py3-none-any.whl

atdata/.gitignore

@@ -0,0 +1 @@
+!manifest/

atdata/__init__.py

@@ -44,6 +44,7 @@ from .dataset import (
     SampleBatch as SampleBatch,
     Dataset as Dataset,
     packable as packable,
+    write_samples as write_samples,
 )
 
 from .lens import (
@@ -55,6 +56,8 @@ from .lens import (
 from ._hf_api import (
     load_dataset as load_dataset,
     DatasetDict as DatasetDict,
+    get_default_index as get_default_index,
+    set_default_index as set_default_index,
 )
 
 from ._protocols import (
@@ -71,10 +74,37 @@ from ._sources import (
     BlobSource as BlobSource,
 )
 
+from ._exceptions import (
+    AtdataError as AtdataError,
+    LensNotFoundError as LensNotFoundError,
+    SchemaError as SchemaError,
+    SampleKeyError as SampleKeyError,
+    ShardError as ShardError,
+    PartialFailureError as PartialFailureError,
+)
+
 from ._schema_codec import (
     schema_to_type as schema_to_type,
 )
 
+from ._logging import (
+    configure_logging as configure_logging,
+    get_logger as get_logger,
+)
+
+from .repository import (
+    Repository as Repository,
+    create_repository as create_repository,
+)
+
+from .index import (
+    Index as Index,
+)
+
+from .stores import (
+    LocalDiskStore as LocalDiskStore,
+)
+
 from ._cid import (
     generate_cid as generate_cid,
     verify_cid as verify_cid,
@@ -84,6 +114,15 @@ from .promote import (
     promote_to_atmosphere as promote_to_atmosphere,
 )
 
+from .manifest import (
+    ManifestField as ManifestField,
+    ManifestBuilder as ManifestBuilder,
+    ShardManifest as ShardManifest,
+    ManifestWriter as ManifestWriter,
+    QueryExecutor as QueryExecutor,
+    SampleLocation as SampleLocation,
+)
+
 # ATProto integration (lazy import to avoid requiring atproto package)
 from . import atmosphere as atmosphere
 

atdata/_cid.py

@@ -116,29 +116,8 @@ def verify_cid(cid: str, data: Any) -> bool:
     return cid == expected_cid
 
 
-def parse_cid(cid: str) -> dict:
-    """Parse a CID string into its components.
-
-    Args:
-        cid: CID string to parse.
-
-    Returns:
-        Dictionary with 'version', 'codec', and 'hash' keys.
-        The 'hash' value is itself a dict with 'code', 'size', and 'digest'.
-
-    Examples:
-        >>> info = parse_cid('bafyrei...')
-        >>> info['version']
-        1
-        >>> info['codec']
-        113  # 0x71 = dag-cbor
-    """
-    return libipld.decode_cid(cid)
-
-
 __all__ = [
     "generate_cid",
     "generate_cid_from_bytes",
     "verify_cid",
-    "parse_cid",
 ]

atdata/_exceptions.py

@@ -0,0 +1,168 @@
+"""Custom exception hierarchy for atdata.
+
+Provides actionable error messages with contextual help, available
+alternatives, and suggested fix code snippets.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Type
+
+
+class AtdataError(Exception):
+    """Base exception for all atdata errors."""
+
+
+class LensNotFoundError(AtdataError, ValueError):
+    """No lens registered to transform between two sample types.
+
+    Attributes:
+        source_type: The source sample type.
+        view_type: The target view type.
+        available_targets: Types reachable from the source via registered lenses.
+    """
+
+    def __init__(
+        self,
+        source_type: Type,
+        view_type: Type,
+        available_targets: list[tuple[Type, str]] | None = None,
+    ) -> None:
+        self.source_type = source_type
+        self.view_type = view_type
+        self.available_targets = available_targets or []
+
+        src_name = source_type.__name__
+        view_name = view_type.__name__
+
+        lines = [f"No lens transforms {src_name} \u2192 {view_name}"]
+
+        if self.available_targets:
+            lines.append("")
+            lines.append(f"Available lenses from {src_name}:")
+            for target_type, lens_name in self.available_targets:
+                lines.append(
+                    f"  - {src_name} \u2192 {target_type.__name__} (via {lens_name})"
+                )
+
+        lines.append("")
+        lines.append("Did you mean to define:")
+        lines.append("  @lens")
+        lines.append(
+            f"  def {src_name.lower()}_to_{view_name.lower()}(source: {src_name}) -> {view_name}:"
+        )
+        lines.append(f"      return {view_name}(...)")
+
+        super().__init__("\n".join(lines))
+
+
+class SchemaError(AtdataError):
+    """Schema mismatch during sample deserialization.
+
+    Raised when the data in a shard doesn't match the expected sample type.
+
+    Attributes:
+        expected_fields: Fields expected by the sample type.
+        actual_fields: Fields found in the data.
+        sample_type_name: Name of the target sample type.
+    """
+
+    def __init__(
+        self,
+        sample_type_name: str,
+        expected_fields: list[str],
+        actual_fields: list[str],
+    ) -> None:
+        self.sample_type_name = sample_type_name
+        self.expected_fields = expected_fields
+        self.actual_fields = actual_fields
+
+        missing = sorted(set(expected_fields) - set(actual_fields))
+        extra = sorted(set(actual_fields) - set(expected_fields))
+
+        lines = [f"Schema mismatch for {sample_type_name}"]
+        if missing:
+            lines.append(f"  Missing fields: {', '.join(missing)}")
+        if extra:
+            lines.append(f"  Unexpected fields: {', '.join(extra)}")
+        lines.append("")
+        lines.append(f"Expected: {', '.join(sorted(expected_fields))}")
+        lines.append(f"Got:      {', '.join(sorted(actual_fields))}")
+
+        super().__init__("\n".join(lines))
+
+
+class SampleKeyError(AtdataError, KeyError):
+    """Sample with the given key was not found in the dataset.
+
+    Attributes:
+        key: The key that was not found.
+    """
+
+    def __init__(self, key: str) -> None:
+        self.key = key
+        super().__init__(
+            f"Sample with key '{key}' not found in dataset. "
+            f"Note: key lookup requires scanning all shards and is O(n)."
+        )
+
+
+class ShardError(AtdataError):
+    """Error accessing or reading a dataset shard.
+
+    Attributes:
+        shard_id: Identifier of the shard that failed.
+        reason: Human-readable description of what went wrong.
+    """
+
+    def __init__(self, shard_id: str, reason: str) -> None:
+        self.shard_id = shard_id
+        self.reason = reason
+        super().__init__(f"Failed to read shard '{shard_id}': {reason}")
+
+
+class PartialFailureError(AtdataError):
+    """Some shards succeeded but others failed during processing.
+
+    Raised by :meth:`Dataset.process_shards` when at least one shard fails.
+    Provides access to both the successful results and the per-shard errors,
+    enabling retry of only the failed shards.
+
+    Attributes:
+        succeeded_shards: List of shard identifiers that succeeded.
+        failed_shards: List of shard identifiers that failed.
+        errors: Mapping from shard identifier to the exception that occurred.
+        results: Mapping from shard identifier to the result for succeeded shards.
+    """
+
+    def __init__(
+        self,
+        succeeded_shards: list[str],
+        failed_shards: list[str],
+        errors: dict[str, Exception],
+        results: dict[str, object],
+    ) -> None:
+        self.succeeded_shards = succeeded_shards
+        self.failed_shards = failed_shards
+        self.errors = errors
+        self.results = results
+
+        n_ok = len(succeeded_shards)
+        n_fail = len(failed_shards)
+        total = n_ok + n_fail
+
+        lines = [f"{n_fail}/{total} shards failed during processing"]
+        for shard_id in failed_shards[:5]:
+            lines.append(f"  {shard_id}: {errors[shard_id]}")
+        if n_fail > 5:
+            lines.append(f"  ... and {n_fail - 5} more")
+        lines.append("")
+        lines.append(
+            f"Access .succeeded_shards ({n_ok}) and .failed_shards ({n_fail}) "
+            f"to inspect or retry."
+        )
+
+        super().__init__("\n".join(lines))

atdata/_helpers.py

@@ -1,8 +1,7 @@
 """Helper utilities for numpy array serialization.
 
 This module provides utility functions for converting numpy arrays to and from
-bytes for msgpack serialization. The functions use numpy's native save/load
-format to preserve array dtype and shape information.
+bytes for msgpack serialization.
 
 Functions:
     - ``array_to_bytes()``: Serialize numpy array to bytes
@@ -15,10 +14,14 @@ handling of NDArray fields during msgpack packing/unpacking.
 ##
 # Imports
 
+import struct
 from io import BytesIO
 
 import numpy as np
 
+# .npy format magic prefix (used for backward-compatible deserialization)
+_NPY_MAGIC = b"\x93NUMPY"
+
 
 ##
 
@@ -26,35 +29,58 @@ import numpy as np
 def array_to_bytes(x: np.ndarray) -> bytes:
     """Convert a numpy array to bytes for msgpack serialization.
 
-    Uses numpy's native ``save()`` format to preserve array dtype and shape.
+    Uses a compact binary format: a short header (dtype + shape) followed by
+    raw array bytes via ``ndarray.tobytes()``. Falls back to numpy's ``.npy``
+    format for object dtypes that cannot be represented as raw bytes.
 
     Args:
         x: A numpy array to serialize.
 
     Returns:
         Raw bytes representing the serialized array.
-
-    Note:
-        Uses ``allow_pickle=True`` to support object dtypes.
     """
-    np_bytes = BytesIO()
-    np.save(np_bytes, x, allow_pickle=True)
-    return np_bytes.getvalue()
+    if x.dtype == object:
+        buf = BytesIO()
+        np.save(buf, x, allow_pickle=True)
+        return buf.getvalue()
+
+    dtype_str = x.dtype.str.encode()  # e.g. b'<f4'
+    header = struct.pack(f"<B{len(x.shape)}q", len(x.shape), *x.shape)
+    return struct.pack("<B", len(dtype_str)) + dtype_str + header + x.tobytes()
 
 
 def bytes_to_array(b: bytes) -> np.ndarray:
     """Convert serialized bytes back to a numpy array.
 
-    Reverses the serialization performed by ``array_to_bytes()``.
+    Transparently handles both the compact format produced by the current
+    ``array_to_bytes()`` and the legacy ``.npy`` format.
 
     Args:
         b: Raw bytes from a serialized numpy array.
 
     Returns:
         The deserialized numpy array with original dtype and shape.
-
-    Note:
-        Uses ``allow_pickle=True`` to support object dtypes.
     """
-    np_bytes = BytesIO(b)
-    return np.load(np_bytes, allow_pickle=True)
+    if b[:6] == _NPY_MAGIC:
+        return np.load(BytesIO(b), allow_pickle=True)
+
+    # Compact format: dtype_len(1B) + dtype_str + ndim(1B) + shape(ndim×8B) + data
+    if len(b) < 2:
+        raise ValueError(f"Array buffer too short ({len(b)} bytes): need at least 2")
+    dlen = b[0]
+    min_header = 2 + dlen  # dtype_len + dtype_str + ndim
+    if len(b) < min_header:
+        raise ValueError(
+            f"Array buffer too short ({len(b)} bytes): need at least {min_header} for header"
+        )
+    dtype = np.dtype(b[1 : 1 + dlen].decode())
+    ndim = b[1 + dlen]
+    offset = 2 + dlen
+    min_with_shape = offset + ndim * 8
+    if len(b) < min_with_shape:
+        raise ValueError(
+            f"Array buffer too short ({len(b)} bytes): need at least {min_with_shape} for shape"
+        )
+    shape = struct.unpack_from(f"<{ndim}q", b, offset)
+    offset += ndim * 8
+    return np.frombuffer(b, dtype=dtype, offset=offset).reshape(shape).copy()

atdata/_hf_api.py

@@ -29,8 +29,10 @@ Examples:
 from __future__ import annotations
 
 import re
+import threading
 from pathlib import Path
 from typing import (
+    Any,
     TYPE_CHECKING,
     Generic,
     Mapping,
@@ -40,9 +42,9 @@ from typing import (
     overload,
 )
 
-from .dataset import Dataset, PackableSample, DictSample
+from .dataset import Dataset, DictSample
 from ._sources import URLSource, S3Source
-from ._protocols import DataSource
+from ._protocols import DataSource, Packable
 
 if TYPE_CHECKING:
     from ._protocols import AbstractIndex
@@ -50,7 +52,60 @@ if TYPE_CHECKING:
 ##
 # Type variables
 
-ST = TypeVar("ST", bound=PackableSample)
+ST = TypeVar("ST", bound=Packable)
+
+
+##
+# Default Index singleton
+
+_default_index: "Index | None" = None  # noqa: F821 (forward ref)
+_default_index_lock = threading.Lock()
+
+
+def get_default_index() -> "Index":  # noqa: F821
+    """Get or create the module-level default Index.
+
+    The default Index uses Redis for local storage (backwards-compatible
+    default) and an anonymous Atmosphere for read-only public data
+    resolution.
+
+    The default is created lazily on first access and cached for the
+    lifetime of the process.
+
+    Returns:
+        The default Index instance.
+
+    Examples:
+        >>> index = get_default_index()
+        >>> entry = index.get_dataset("local/mnist")
+    """
+    global _default_index
+    if _default_index is None:
+        with _default_index_lock:
+            if _default_index is None:
+                from .local import Index
+
+                _default_index = Index()
+    return _default_index
+
+
+def set_default_index(index: "Index") -> None:  # noqa: F821
+    """Override the module-level default Index.
+
+    Use this to configure a custom default Index with specific repositories,
+    an authenticated atmosphere client, or non-default providers.
+
+    Args:
+        index: The Index instance to use as the default.
+
+    Examples:
+        >>> from atdata.local import Index
+        >>> from atdata.providers import create_provider
+        >>> custom = Index(provider=create_provider("sqlite"))
+        >>> set_default_index(custom)
+    """
+    global _default_index
+    _default_index = index
 
 
 ##
@@ -74,10 +129,11 @@ class DatasetDict(Generic[ST], dict):
         >>>
         >>> # Iterate over all splits
         >>> for split_name, dataset in ds_dict.items():
-        ...     print(f"{split_name}: {len(dataset.shard_list)} shards")
+        ...     print(f"{split_name}: {len(dataset.list_shards())} shards")
     """
 
-    # TODO The above has a line for "Parameters:" that should be "Type Parameters:"; this is a temporary fix for `quartodoc` auto-generation bugs.
+    # Note: The docstring uses "Parameters:" for type parameters as a workaround
+    # for quartodoc not supporting "Type Parameters:" sections.
 
     def __init__(
         self,
@@ -134,6 +190,37 @@ class DatasetDict(Generic[ST], dict):
         """
         return {name: len(ds.list_shards()) for name, ds in self.items()}
 
+    # Methods proxied to the sole Dataset when only one split exists.
+    _DATASET_METHODS = frozenset(
+        {
+            "ordered",
+            "shuffled",
+            "as_type",
+            "list_shards",
+            "head",
+        }
+    )
+
+    def __getattr__(self, name: str) -> Any:
+        """Proxy common Dataset methods when this dict has exactly one split.
+
+        When a ``DatasetDict`` contains a single split, calling iteration
+        methods like ``.ordered()`` or ``.shuffled()`` is forwarded to the
+        contained ``Dataset`` for convenience.  Multi-split dicts raise
+        ``AttributeError`` with a hint to select a split explicitly.
+        """
+        if name in self._DATASET_METHODS:
+            if len(self) == 1:
+                return getattr(next(iter(self.values())), name)
+            splits = ", ".join(f"'{k}'" for k in self.keys())
+            raise AttributeError(
+                f"'{type(self).__name__}' has {len(self)} splits ({splits}). "
+                f"Select one first, e.g. ds_dict['{next(iter(self.keys()))}'].{name}()"
+            )
+        raise AttributeError(
+            f"'{type(self).__name__}' object has no attribute '{name}'"
+        )
+
 
 ##
 # Path resolution utilities
@@ -459,7 +546,7 @@ def _resolve_indexed_path(
     handle_or_did, dataset_name = _parse_indexed_path(path)
 
     # For AtmosphereIndex, we need to resolve handle to DID first
-    # For LocalIndex, the handle is ignored and we just look up by name
+    # For local Index, the handle is ignored and we just look up by name
     entry = index.get_dataset(dataset_name)
     data_urls = entry.data_urls
 
@@ -624,16 +711,13 @@ def load_dataset(
         >>> train_ds = load_dataset("./data/train-*.tar", TextData, split="train")
         >>>
         >>> # Load from index with auto-type resolution
-        >>> index = LocalIndex()
+        >>> index = Index()
         >>> ds = load_dataset("@local/my-dataset", index=index, split="train")
     """
     # Handle @handle/dataset indexed path resolution
     if _is_indexed_path(path):
         if index is None:
-            raise ValueError(
-                f"Index required for indexed path: {path}. "
-                "Pass index=LocalIndex() or index=AtmosphereIndex(client)."
-            )
+            index = get_default_index()
 
         source, schema_ref = _resolve_indexed_path(path, index)
 

atdata/_logging.py

@@ -0,0 +1,70 @@
+"""Pluggable logging for atdata.
+
+Provides a thin abstraction over Python's stdlib ``logging`` module that can
+be replaced with ``structlog`` or any other logger implementing the standard
+``debug``/``info``/``warning``/``error`` interface.
+
+Usage::
+
+    # Default: stdlib logging (no config needed)
+    from atdata._logging import get_logger
+    log = get_logger()
+    log.info("processing shard", extra={"shard": "data-000.tar"})
+
+    # Plug in structlog (or any compatible logger):
+    import structlog
+    import atdata
+    atdata.configure_logging(structlog.get_logger())
+
+The module also exports a lightweight ``LoggerProtocol`` for type checking
+custom logger implementations.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class LoggerProtocol(Protocol):
+    """Minimal interface that a pluggable logger must satisfy."""
+
+    def debug(self, msg: str, *args: Any, **kwargs: Any) -> None: ...
+    def info(self, msg: str, *args: Any, **kwargs: Any) -> None: ...
+    def warning(self, msg: str, *args: Any, **kwargs: Any) -> None: ...
+    def error(self, msg: str, *args: Any, **kwargs: Any) -> None: ...
+
+
+# ---------------------------------------------------------------------------
+# Module-level state
+# ---------------------------------------------------------------------------
+
+_logger: LoggerProtocol = logging.getLogger("atdata")
+
+
+def configure_logging(logger: LoggerProtocol) -> None:
+    """Replace the default logger with a custom implementation.
+
+    The provided logger must implement ``debug``, ``info``, ``warning``, and
+    ``error`` methods. Both ``structlog`` bound loggers and stdlib
+    ``logging.Logger`` instances satisfy this interface.
+
+    Args:
+        logger: A logger instance implementing :class:`LoggerProtocol`.
+
+    Examples:
+        >>> import structlog
+        >>> atdata.configure_logging(structlog.get_logger())
+    """
+    global _logger
+    _logger = logger
+
+
+def get_logger() -> LoggerProtocol:
+    """Return the currently configured logger.
+
+    Returns the stdlib ``logging.getLogger("atdata")`` by default, or
+    whatever was last set via :func:`configure_logging`.
+    """
+    return _logger