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

atdata/.gitignore

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

atdata/__init__.py

@@ -55,6 +55,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 +73,29 @@ 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 ._cid import (
     generate_cid as generate_cid,
     verify_cid as verify_cid,
@@ -84,8 +105,17 @@ 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
 
 # CLI entry point
-from .cli import main as main
+from .cli import main as main

atdata/_cid.py

@@ -12,13 +12,11 @@ The CIDs generated here use:
 This ensures compatibility with ATProto's CID requirements and enables
 seamless promotion from local storage to atmosphere (ATProto network).
 
-Example:
-    ::
-
-        >>> schema = {"name": "ImageSample", "version": "1.0.0", "fields": [...]}
-        >>> cid = generate_cid(schema)
-        >>> print(cid)
-        bafyreihffx5a2e7k6r5zqgp5iwpjqr2gfyheqhzqtlxagvqjqyxzqpzqaa
+Examples:
+    >>> schema = {"name": "ImageSample", "version": "1.0.0", "fields": [...]}
+    >>> cid = generate_cid(schema)
+    >>> print(cid)
+    bafyreihffx5a2e7k6r5zqgp5iwpjqr2gfyheqhzqtlxagvqjqyxzqpzqaa
 """
 
 import hashlib
@@ -50,11 +48,9 @@ def generate_cid(data: Any) -> str:
     Raises:
         ValueError: If the data cannot be encoded as DAG-CBOR.
 
-    Example:
-        ::
-
-            >>> generate_cid({"name": "test", "value": 42})
-            'bafyrei...'
+    Examples:
+        >>> generate_cid({"name": "test", "value": 42})
+        'bafyrei...'
     """
     # Encode data as DAG-CBOR
     try:
@@ -68,7 +64,9 @@ def generate_cid(data: Any) -> str:
     # Build raw CID bytes:
     # CIDv1 = version(1) + codec(dag-cbor) + multihash
     # Multihash = code(sha256) + size(32) + digest
-    raw_cid_bytes = bytes([CID_VERSION_1, CODEC_DAG_CBOR, HASH_SHA256, SHA256_SIZE]) + sha256_hash
+    raw_cid_bytes = (
+        bytes([CID_VERSION_1, CODEC_DAG_CBOR, HASH_SHA256, SHA256_SIZE]) + sha256_hash
+    )
 
     # Encode to base32 multibase string
     return libipld.encode_cid(raw_cid_bytes)
@@ -86,14 +84,14 @@ def generate_cid_from_bytes(data_bytes: bytes) -> str:
     Returns:
         CIDv1 string in base32 multibase format.
 
-    Example:
-        ::
-
-            >>> cbor_bytes = libipld.encode_dag_cbor({"key": "value"})
-            >>> cid = generate_cid_from_bytes(cbor_bytes)
+    Examples:
+        >>> cbor_bytes = libipld.encode_dag_cbor({"key": "value"})
+        >>> cid = generate_cid_from_bytes(cbor_bytes)
     """
     sha256_hash = hashlib.sha256(data_bytes).digest()
-    raw_cid_bytes = bytes([CID_VERSION_1, CODEC_DAG_CBOR, HASH_SHA256, SHA256_SIZE]) + sha256_hash
+    raw_cid_bytes = (
+        bytes([CID_VERSION_1, CODEC_DAG_CBOR, HASH_SHA256, SHA256_SIZE]) + sha256_hash
+    )
     return libipld.encode_cid(raw_cid_bytes)
 
 
@@ -107,14 +105,12 @@ def verify_cid(cid: str, data: Any) -> bool:
     Returns:
         True if the CID matches the data, False otherwise.
 
-    Example:
-        ::
-
-            >>> cid = generate_cid({"name": "test"})
-            >>> verify_cid(cid, {"name": "test"})
-            True
-            >>> verify_cid(cid, {"name": "different"})
-            False
+    Examples:
+        >>> cid = generate_cid({"name": "test"})
+        >>> verify_cid(cid, {"name": "test"})
+        True
+        >>> verify_cid(cid, {"name": "different"})
+        False
     """
     expected_cid = generate_cid(data)
     return cid == expected_cid
@@ -130,14 +126,12 @@ def parse_cid(cid: str) -> dict:
         Dictionary with 'version', 'codec', and 'hash' keys.
         The 'hash' value is itself a dict with 'code', 'size', and 'digest'.
 
-    Example:
-        ::
-
-            >>> info = parse_cid('bafyrei...')
-            >>> info['version']
-            1
-            >>> info['codec']
-            113  # 0x71 = dag-cbor
+    Examples:
+        >>> info = parse_cid('bafyrei...')
+        >>> info['version']
+        1
+        >>> info['codec']
+        113  # 0x71 = dag-cbor
     """
     return libipld.decode_cid(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,44 +14,61 @@ 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"
+
 
 ##
 
-def array_to_bytes( x: np.ndarray ) -> bytes:
+
+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()
 
-def bytes_to_array( b: bytes ) -> np.ndarray:
+    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
+    dlen = b[0]
+    dtype = np.dtype(b[1 : 1 + dlen].decode())
+    ndim = b[1 + dlen]
+    offset = 2 + dlen
+    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

@@ -9,28 +9,27 @@ Key differences from HuggingFace Datasets:
 - Built on WebDataset for efficient streaming of large datasets
 - No Arrow caching layer (WebDataset handles remote/local transparently)
 
-Example:
-    ::
-
-        >>> import atdata
-        >>> from atdata import load_dataset
-        >>>
-        >>> @atdata.packable
-        ... class MyData:
-        ...     text: str
-        ...     label: int
-        >>>
-        >>> # Load a single split
-        >>> ds = load_dataset("path/to/train-{000000..000099}.tar", MyData, split="train")
-        >>>
-        >>> # Load all splits (returns DatasetDict)
-        >>> ds_dict = load_dataset("path/to/{train,test}-*.tar", MyData)
-        >>> train_ds = ds_dict["train"]
+Examples:
+    >>> import atdata
+    >>> from atdata import load_dataset
+    >>>
+    >>> @atdata.packable
+    ... class MyData:
+    ...     text: str
+    ...     label: int
+    >>>
+    >>> # Load a single split
+    >>> ds = load_dataset("path/to/train-{000000..000099}.tar", MyData, split="train")
+    >>>
+    >>> # Load all splits (returns DatasetDict)
+    >>> ds_dict = load_dataset("path/to/{train,test}-*.tar", MyData)
+    >>> train_ds = ds_dict["train"]
 """
 
 from __future__ import annotations
 
 import re
+import threading
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
@@ -42,18 +41,70 @@ 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
-    from .local import S3DataStore
 
 ##
 # 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 AtmosphereClient 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
 
 
 ##
@@ -70,18 +121,18 @@ class DatasetDict(Generic[ST], dict):
     Parameters:
         ST: The sample type for all datasets in this dict.
 
-    Example:
-        ::
-
-            >>> ds_dict = load_dataset("path/to/data", MyData)
-            >>> train = ds_dict["train"]
-            >>> test = ds_dict["test"]
-            >>>
-            >>> # Iterate over all splits
-            >>> for split_name, dataset in ds_dict.items():
-            ...     print(f"{split_name}: {len(dataset.shard_list)} shards")
+    Examples:
+        >>> ds_dict = load_dataset("path/to/data", MyData)
+        >>> train = ds_dict["train"]
+        >>> test = ds_dict["test"]
+        >>>
+        >>> # Iterate over all splits
+        >>> for split_name, dataset in ds_dict.items():
+        ...     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,
@@ -463,12 +514,12 @@ 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
 
     # Check if index has a data store
-    if hasattr(index, 'data_store') and index.data_store is not None:
+    if hasattr(index, "data_store") and index.data_store is not None:
         store = index.data_store
 
         # Import here to avoid circular imports at module level
@@ -613,38 +664,35 @@ def load_dataset(
         FileNotFoundError: If no data files are found at the path.
         KeyError: If dataset not found in index.
 
-    Example:
-        ::
-
-            >>> # Load without type - get DictSample for exploration
-            >>> ds = load_dataset("./data/train.tar", split="train")
-            >>> for sample in ds.ordered():
-            ...     print(sample.keys())  # Explore fields
-            ...     print(sample["text"]) # Dict-style access
-            ...     print(sample.label)   # Attribute access
-            >>>
-            >>> # Convert to typed schema
-            >>> typed_ds = ds.as_type(TextData)
-            >>>
-            >>> # Or load with explicit type directly
-            >>> train_ds = load_dataset("./data/train-*.tar", TextData, split="train")
-            >>>
-            >>> # Load from index with auto-type resolution
-            >>> index = LocalIndex()
-            >>> ds = load_dataset("@local/my-dataset", index=index, split="train")
+    Examples:
+        >>> # Load without type - get DictSample for exploration
+        >>> ds = load_dataset("./data/train.tar", split="train")
+        >>> for sample in ds.ordered():
+        ...     print(sample.keys())  # Explore fields
+        ...     print(sample["text"]) # Dict-style access
+        ...     print(sample.label)   # Attribute access
+        >>>
+        >>> # Convert to typed schema
+        >>> typed_ds = ds.as_type(TextData)
+        >>>
+        >>> # Or load with explicit type directly
+        >>> train_ds = load_dataset("./data/train-*.tar", TextData, split="train")
+        >>>
+        >>> # Load from index with auto-type resolution
+        >>> 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)
 
         # Resolve sample_type from schema if not provided
-        resolved_type: Type = sample_type if sample_type is not None else index.decode_schema(schema_ref)
+        resolved_type: Type = (
+            sample_type if sample_type is not None else index.decode_schema(schema_ref)
+        )
 
         # Create dataset from the resolved source (includes credentials if S3)
         ds = Dataset[resolved_type](source)
@@ -653,7 +701,9 @@ def load_dataset(
             # Indexed datasets are single-split by default
             return ds
 
-        return DatasetDict({"train": ds}, sample_type=resolved_type, streaming=streaming)
+        return DatasetDict(
+            {"train": ds}, sample_type=resolved_type, streaming=streaming
+        )
 
     # Use DictSample as default when no type specified
     resolved_type = sample_type if sample_type is not None else DictSample