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

atdata/_sources.py

@@ -13,22 +13,20 @@ The key insight is that WebDataset's tar_file_expander only needs
 By providing streams directly, we can support private repos, custom
 endpoints, and future backends like ATProto blobs.
 
-Example:
-    ::
-
-        >>> # Standard URL (uses WebDataset's gopen)
-        >>> source = URLSource("https://example.com/data-{000..009}.tar")
-        >>> ds = Dataset[MySample](source)
-        >>>
-        >>> # Private S3 with credentials
-        >>> source = S3Source(
-        ...     bucket="my-bucket",
-        ...     keys=["train/shard-000.tar", "train/shard-001.tar"],
-        ...     endpoint="https://my-r2.cloudflarestorage.com",
-        ...     access_key="...",
-        ...     secret_key="...",
-        ... )
-        >>> ds = Dataset[MySample](source)
+Examples:
+    >>> # Standard URL (uses WebDataset's gopen)
+    >>> source = URLSource("https://example.com/data-{000..009}.tar")
+    >>> ds = Dataset[MySample](source)
+    >>>
+    >>> # Private S3 with credentials
+    >>> source = S3Source(
+    ...     bucket="my-bucket",
+    ...     keys=["train/shard-000.tar", "train/shard-001.tar"],
+    ...     endpoint="https://my-r2.cloudflarestorage.com",
+    ...     access_key="...",
+    ...     secret_key="...",
+    ... )
+    >>> ds = Dataset[MySample](source)
 """
 
 from __future__ import annotations
@@ -54,12 +52,10 @@ class URLSource:
     Attributes:
         url: URL or brace pattern for the shards.
 
-    Example:
-        ::
-
-            >>> source = URLSource("https://example.com/train-{000..009}.tar")
-            >>> for shard_id, stream in source.shards:
-            ...     print(f"Streaming {shard_id}")
+    Examples:
+        >>> source = URLSource("https://example.com/train-{000..009}.tar")
+        >>> for shard_id, stream in source.shards:
+        ...     print(f"Streaming {shard_id}")
     """
 
     url: str
@@ -131,18 +127,16 @@ class S3Source:
         secret_key: Optional AWS secret access key.
         region: Optional AWS region (defaults to us-east-1).
 
-    Example:
-        ::
-
-            >>> source = S3Source(
-            ...     bucket="my-datasets",
-            ...     keys=["train/shard-000.tar", "train/shard-001.tar"],
-            ...     endpoint="https://abc123.r2.cloudflarestorage.com",
-            ...     access_key="AKIAIOSFODNN7EXAMPLE",
-            ...     secret_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
-            ... )
-            >>> for shard_id, stream in source.shards:
-            ...     process(stream)
+    Examples:
+        >>> source = S3Source(
+        ...     bucket="my-datasets",
+        ...     keys=["train/shard-000.tar", "train/shard-001.tar"],
+        ...     endpoint="https://abc123.r2.cloudflarestorage.com",
+        ...     access_key="AKIAIOSFODNN7EXAMPLE",
+        ...     secret_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
+        ... )
+        >>> for shard_id, stream in source.shards:
+        ...     process(stream)
     """
 
     bucket: str
@@ -173,7 +167,9 @@ class S3Source:
             client_kwargs["region_name"] = self.region
         elif not self.endpoint:
             # Default region for AWS S3
-            client_kwargs["region_name"] = os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
+            client_kwargs["region_name"] = os.environ.get(
+                "AWS_DEFAULT_REGION", "us-east-1"
+            )
 
         self._client = boto3.client("s3", **client_kwargs)
         return self._client
@@ -225,7 +221,7 @@ class S3Source:
         if not shard_id.startswith(f"s3://{self.bucket}/"):
             raise KeyError(f"Shard not in this bucket: {shard_id}")
 
-        key = shard_id[len(f"s3://{self.bucket}/"):]
+        key = shard_id[len(f"s3://{self.bucket}/") :]
         client = self._get_client()
         response = client.get_object(Bucket=self.bucket, Key=key)
         return response["Body"]
@@ -258,13 +254,11 @@ class S3Source:
         Raises:
             ValueError: If URLs are not valid s3:// URLs or span multiple buckets.
 
-        Example:
-            ::
-
-                >>> source = S3Source.from_urls(
-                ...     ["s3://my-bucket/train-000.tar", "s3://my-bucket/train-001.tar"],
-                ...     endpoint="https://r2.example.com",
-                ... )
+        Examples:
+            >>> source = S3Source.from_urls(
+            ...     ["s3://my-bucket/train-000.tar", "s3://my-bucket/train-001.tar"],
+            ...     endpoint="https://r2.example.com",
+            ... )
         """
         if not urls:
             raise ValueError("urls cannot be empty")
@@ -317,15 +311,13 @@ class S3Source:
         Returns:
             Configured S3Source.
 
-        Example:
-            ::
-
-                >>> creds = {
-                ...     "AWS_ACCESS_KEY_ID": "...",
-                ...     "AWS_SECRET_ACCESS_KEY": "...",
-                ...     "AWS_ENDPOINT": "https://r2.example.com",
-                ... }
-                >>> source = S3Source.from_credentials(creds, "my-bucket", ["data.tar"])
+        Examples:
+            >>> creds = {
+            ...     "AWS_ACCESS_KEY_ID": "...",
+            ...     "AWS_SECRET_ACCESS_KEY": "...",
+            ...     "AWS_ENDPOINT": "https://r2.example.com",
+            ... }
+            >>> source = S3Source.from_credentials(creds, "my-bucket", ["data.tar"])
         """
         return cls(
             bucket=bucket,
@@ -352,22 +344,22 @@ class BlobSource:
         blob_refs: List of blob reference dicts with 'did' and 'cid' keys.
         pds_endpoint: Optional PDS endpoint URL. If not provided, resolved from DID.
 
-    Example:
-        ::
-
-            >>> source = BlobSource(
-            ...     blob_refs=[
-            ...         {"did": "did:plc:abc123", "cid": "bafyrei..."},
-            ...         {"did": "did:plc:abc123", "cid": "bafyrei..."},
-            ...     ],
-            ... )
-            >>> for shard_id, stream in source.shards:
-            ...     process(stream)
+    Examples:
+        >>> source = BlobSource(
+        ...     blob_refs=[
+        ...         {"did": "did:plc:abc123", "cid": "bafyrei..."},
+        ...         {"did": "did:plc:abc123", "cid": "bafyrei..."},
+        ...     ],
+        ... )
+        >>> for shard_id, stream in source.shards:
+        ...     process(stream)
     """
 
     blob_refs: list[dict[str, str]]
     pds_endpoint: str | None = None
-    _endpoint_cache: dict[str, str] = field(default_factory=dict, repr=False, compare=False)
+    _endpoint_cache: dict[str, str] = field(
+        default_factory=dict, repr=False, compare=False
+    )
 
     def _resolve_pds_endpoint(self, did: str) -> str:
         """Resolve PDS endpoint for a DID, with caching."""
@@ -459,6 +451,7 @@ class BlobSource:
         url = self._get_blob_url(did, cid)
 
         import requests
+
         response = requests.get(url, stream=True, timeout=60)
         response.raise_for_status()
         return response.raw

atdata/_stub_manager.py

@@ -8,20 +8,18 @@ Unlike simple .pyi stubs, the generated modules are actual Python code that
 can be imported at runtime. This allows ``decode_schema`` to return properly
 typed classes that work with both static type checkers and runtime.
 
-Example:
-    ::
-
-        >>> from atdata.local import Index
-        >>>
-        >>> # Enable auto-stub generation
-        >>> index = Index(auto_stubs=True)
-        >>>
-        >>> # Modules are generated automatically on decode_schema
-        >>> MyType = index.decode_schema("atdata://local/sampleSchema/MySample@1.0.0")
-        >>> # MyType is now properly typed for IDE autocomplete!
-        >>>
-        >>> # Get the stub directory path for IDE configuration
-        >>> print(f"Add to IDE: {index.stub_dir}")
+Examples:
+    >>> from atdata.local import Index
+    >>>
+    >>> # Enable auto-stub generation
+    >>> index = Index(auto_stubs=True)
+    >>>
+    >>> # Modules are generated automatically on decode_schema
+    >>> MyType = index.decode_schema("atdata://local/sampleSchema/MySample@1.0.0")
+    >>> # MyType is now properly typed for IDE autocomplete!
+    >>>
+    >>> # Get the stub directory path for IDE configuration
+    >>> print(f"Add to IDE: {index.stub_dir}")
 """
 
 from pathlib import Path
@@ -101,14 +99,12 @@ class StubManager:
     Args:
         stub_dir: Directory to write module files. Defaults to ``~/.atdata/stubs/``.
 
-    Example:
-        ::
-
-            >>> manager = StubManager()
-            >>> schema_dict = {"name": "MySample", "version": "1.0.0", "fields": [...]}
-            >>> SampleClass = manager.ensure_module(schema_dict)
-            >>> print(manager.stub_dir)
-            /Users/you/.atdata/stubs
+    Examples:
+        >>> manager = StubManager()
+        >>> schema_dict = {"name": "MySample", "version": "1.0.0", "fields": [...]}
+        >>> SampleClass = manager.ensure_module(schema_dict)
+        >>> print(manager.stub_dir)
+        /Users/you/.atdata/stubs
     """
 
     def __init__(self, stub_dir: Optional[Union[str, Path]] = None):
@@ -157,7 +153,9 @@ class StubManager:
         """Alias for _module_filename for backwards compatibility."""
         return self._module_filename(name, version)
 
-    def _module_path(self, name: str, version: str, authority: str = DEFAULT_AUTHORITY) -> Path:
+    def _module_path(
+        self, name: str, version: str, authority: str = DEFAULT_AUTHORITY
+    ) -> Path:
         """Get full path to module file for a schema.
 
         Args:
@@ -170,7 +168,9 @@ class StubManager:
         """
         return self._stub_dir / authority / self._module_filename(name, version)
 
-    def _stub_path(self, name: str, version: str, authority: str = DEFAULT_AUTHORITY) -> Path:
+    def _stub_path(
+        self, name: str, version: str, authority: str = DEFAULT_AUTHORITY
+    ) -> Path:
         """Alias for _module_path for backwards compatibility."""
         return self._module_path(name, version, authority)
 
@@ -211,7 +211,9 @@ class StubManager:
         authority_dir.mkdir(parents=True, exist_ok=True)
         init_path = authority_dir / "__init__.py"
         if not init_path.exists():
-            init_path.write_text(f'"""Auto-generated schema modules for {authority}."""\n')
+            init_path.write_text(
+                f'"""Auto-generated schema modules for {authority}."""\n'
+            )
 
     def _write_module_atomic(self, path: Path, content: str, authority: str) -> None:
         """Write module file atomically using temp file and rename.
@@ -359,7 +361,9 @@ class StubManager:
 
         return cls
 
-    def _import_class_from_module(self, module_path: Path, class_name: str) -> Optional[Type]:
+    def _import_class_from_module(
+        self, module_path: Path, class_name: str
+    ) -> Optional[Type]:
         """Import a class from a generated module file.
 
         Uses importlib to dynamically load the module and extract the class.
@@ -399,6 +403,7 @@ class StubManager:
     def _print_ide_hint(self) -> None:
         """Print a one-time hint about IDE configuration."""
         import sys as _sys
+
         print(
             f"\n[atdata] Generated schema module in: {self._stub_dir}\n"
             f"[atdata] For IDE support, add this path to your type checker:\n"

atdata/_type_utils.py

@@ -9,15 +9,29 @@ from typing import Any, get_origin, get_args, Union
 
 # Mapping from numpy dtype strings to schema dtype names
 NUMPY_DTYPE_MAP = {
-    "float16": "float16", "float32": "float32", "float64": "float64",
-    "int8": "int8", "int16": "int16", "int32": "int32", "int64": "int64",
-    "uint8": "uint8", "uint16": "uint16", "uint32": "uint32", "uint64": "uint64",
-    "bool": "bool", "complex64": "complex64", "complex128": "complex128",
+    "float16": "float16",
+    "float32": "float32",
+    "float64": "float64",
+    "int8": "int8",
+    "int16": "int16",
+    "int32": "int32",
+    "int64": "int64",
+    "uint8": "uint8",
+    "uint16": "uint16",
+    "uint32": "uint32",
+    "uint64": "uint64",
+    "bool": "bool",
+    "complex64": "complex64",
+    "complex128": "complex128",
 }
 
 # Mapping from Python primitive types to schema type names
 PRIMITIVE_TYPE_MAP = {
-    str: "str", int: "int", float: "float", bool: "bool", bytes: "bytes",
+    str: "str",
+    int: "int",
+    float: "float",
+    bool: "bool",
+    bytes: "bytes",
 }
 
 
@@ -31,9 +45,13 @@ def numpy_dtype_to_string(dtype: Any) -> str:
         Schema dtype string (e.g., "float32", "int64"). Defaults to "float32".
     """
     dtype_str = str(dtype)
-    for key, value in NUMPY_DTYPE_MAP.items():
+    # Exact match first (handles "float32", "int64", etc.)
+    if dtype_str in NUMPY_DTYPE_MAP:
+        return NUMPY_DTYPE_MAP[dtype_str]
+    # Substring match, longest keys first to avoid "int8" matching "uint8"
+    for key in sorted(NUMPY_DTYPE_MAP, key=len, reverse=True):
         if key in dtype_str:
-            return value
+            return NUMPY_DTYPE_MAP[key]
     return "float32"
 
 
@@ -88,3 +106,25 @@ def extract_ndarray_dtype(python_type: Any) -> str:
         if dtype_arg is not None:
             return numpy_dtype_to_string(dtype_arg)
     return "float32"
+
+
+def parse_semver(version: str) -> tuple[int, int, int]:
+    """Parse a semantic version string into a comparable tuple.
+
+    Args:
+        version: A ``"major.minor.patch"`` version string.
+
+    Returns:
+        Tuple of (major, minor, patch) integers.
+
+    Raises:
+        ValueError: If the version string is not valid semver.
+
+    Examples:
+        >>> parse_semver("1.2.3")
+        (1, 2, 3)
+    """
+    parts = version.split(".")
+    if len(parts) != 3:
+        raise ValueError(f"Invalid semver: {version}")
+    return int(parts[0]), int(parts[1]), int(parts[2])

atdata/atmosphere/__init__.py

@@ -15,16 +15,14 @@ The ATProto integration is additive - existing atdata functionality continues
 to work unchanged. These features are opt-in for users who want to publish
 or discover datasets on the ATProto network.
 
-Example:
-    ::
-
-        >>> from atdata.atmosphere import AtmosphereClient, SchemaPublisher
-        >>>
-        >>> client = AtmosphereClient()
-        >>> client.login("handle.bsky.social", "app-password")
-        >>>
-        >>> publisher = SchemaPublisher(client)
-        >>> schema_uri = publisher.publish(MySampleType, version="1.0.0")
+Examples:
+    >>> from atdata.atmosphere import AtmosphereClient, SchemaPublisher
+    >>>
+    >>> client = AtmosphereClient()
+    >>> client.login("handle.bsky.social", "app-password")
+    >>>
+    >>> publisher = SchemaPublisher(client)
+    >>> schema_uri = publisher.publish(MySampleType, version="1.0.0")
 
 Note:
     This module requires the ``atproto`` package to be installed::
@@ -86,6 +84,7 @@ class AtmosphereIndexEntry:
     def metadata(self) -> Optional[dict]:
         """Metadata from the record, if any."""
         import msgpack
+
         metadata_bytes = self._record.get("metadata")
         if metadata_bytes is None:
             return None
@@ -100,25 +99,25 @@ class AtmosphereIndexEntry:
 class AtmosphereIndex:
     """ATProto index implementing AbstractIndex protocol.
 
+    .. deprecated::
+        Use ``atdata.Index(atmosphere=client)`` instead.  ``AtmosphereIndex``
+        is retained for backwards compatibility and will be removed in a
+        future release.
+
     Wraps SchemaPublisher/Loader and DatasetPublisher/Loader to provide
-    a unified interface compatible with LocalIndex.
+    a unified interface compatible with Index.
 
     Optionally accepts a ``PDSBlobStore`` for writing dataset shards as
     ATProto blobs, enabling fully decentralized dataset storage.
 
-    Example:
-        ::
-
-            >>> client = AtmosphereClient()
-            >>> client.login("handle.bsky.social", "app-password")
-            >>>
-            >>> # Without blob storage (external URLs only)
-            >>> index = AtmosphereIndex(client)
-            >>>
-            >>> # With PDS blob storage
-            >>> store = PDSBlobStore(client)
-            >>> index = AtmosphereIndex(client, data_store=store)
-            >>> entry = index.insert_dataset(dataset, name="my-data")
+    Examples:
+        >>> # Preferred: use unified Index
+        >>> from atdata.local import Index
+        >>> from atdata.atmosphere import AtmosphereClient
+        >>> index = Index(atmosphere=client)
+        >>>
+        >>> # Legacy (deprecated)
+        >>> index = AtmosphereIndex(client)
     """
 
     def __init__(
@@ -134,6 +133,14 @@ class AtmosphereIndex:
             data_store: Optional PDSBlobStore for writing shards as blobs.
                 If provided, insert_dataset will upload shards to PDS.
         """
+        import warnings
+
+        warnings.warn(
+            "AtmosphereIndex is deprecated. Use atdata.Index(atmosphere=client) "
+            "instead for unified index access.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         self.client = client
         self._schema_publisher = SchemaPublisher(client)
         self._schema_loader = SchemaLoader(client)

atdata/atmosphere/_types.py

@@ -19,16 +19,14 @@ class AtUri:
 
     AT URIs follow the format: at://<authority>/<collection>/<rkey>
 
-    Example:
-        ::
-
-            >>> uri = AtUri.parse("at://did:plc:abc123/ac.foundation.dataset.sampleSchema/xyz")
-            >>> uri.authority
-            'did:plc:abc123'
-            >>> uri.collection
-            'ac.foundation.dataset.sampleSchema'
-            >>> uri.rkey
-            'xyz'
+    Examples:
+        >>> uri = AtUri.parse("at://did:plc:abc123/ac.foundation.dataset.sampleSchema/xyz")
+        >>> uri.authority
+        'did:plc:abc123'
+        >>> uri.collection
+        'ac.foundation.dataset.sampleSchema'
+        >>> uri.rkey
+        'xyz'
     """
 
     authority: str
@@ -58,7 +56,9 @@ class AtUri:
 
         parts = uri[5:].split("/")
         if len(parts) < 3:
-            raise ValueError(f"Invalid AT URI: expected authority/collection/rkey: {uri}")
+            raise ValueError(
+                f"Invalid AT URI: expected authority/collection/rkey: {uri}"
+            )
 
         return cls(
             authority=parts[0],

atdata/atmosphere/client.py

@@ -18,6 +18,7 @@ def _get_atproto_client_class():
     if _atproto_client_class is None:
         try:
             from atproto import Client
+
             _atproto_client_class = Client
         except ImportError as e:
             raise ImportError(
@@ -33,13 +34,11 @@ class AtmosphereClient:
     This class wraps the atproto SDK client and provides higher-level methods
     for working with atdata records (schemas, datasets, lenses).
 
-    Example:
-        ::
-
-            >>> client = AtmosphereClient()
-            >>> client.login("alice.bsky.social", "app-password")
-            >>> print(client.did)
-            'did:plc:...'
+    Examples:
+        >>> client = AtmosphereClient()
+        >>> client.login("alice.bsky.social", "app-password")
+        >>> print(client.did)
+        'did:plc:...'
 
     Note:
         The password should be an app-specific password, not your main account
@@ -327,7 +326,11 @@ class AtmosphereClient:
         # Convert to dict format suitable for embedding in records
         return {
             "$type": "blob",
-            "ref": {"$link": blob_ref.ref.link if hasattr(blob_ref.ref, "link") else str(blob_ref.ref)},
+            "ref": {
+                "$link": blob_ref.ref.link
+                if hasattr(blob_ref.ref, "link")
+                else str(blob_ref.ref)
+            },
             "mimeType": blob_ref.mime_type,
             "size": blob_ref.size,
         }

atdata/atmosphere/lens.py

@@ -21,6 +21,7 @@ from ._types import (
 
 # Import for type checking only
 from typing import TYPE_CHECKING
+
 if TYPE_CHECKING:
     from ..lens import Lens
 
@@ -31,26 +32,24 @@ class LensPublisher:
     This class creates lens records that reference source and target schemas
     and point to the transformation code in a git repository.
 
-    Example:
-        ::
-
-            >>> @atdata.lens
-            ... def my_lens(source: SourceType) -> TargetType:
-            ...     return TargetType(field=source.other_field)
-            >>>
-            >>> client = AtmosphereClient()
-            >>> client.login("handle", "password")
-            >>>
-            >>> publisher = LensPublisher(client)
-            >>> uri = publisher.publish(
-            ...     name="my_lens",
-            ...     source_schema_uri="at://did:plc:abc/ac.foundation.dataset.sampleSchema/source",
-            ...     target_schema_uri="at://did:plc:abc/ac.foundation.dataset.sampleSchema/target",
-            ...     code_repository="https://github.com/user/repo",
-            ...     code_commit="abc123def456",
-            ...     getter_path="mymodule.lenses:my_lens",
-            ...     putter_path="mymodule.lenses:my_lens_putter",
-            ... )
+    Examples:
+        >>> @atdata.lens
+        ... def my_lens(source: SourceType) -> TargetType:
+        ...     return TargetType(field=source.other_field)
+        >>>
+        >>> client = AtmosphereClient()
+        >>> client.login("handle", "password")
+        >>>
+        >>> publisher = LensPublisher(client)
+        >>> uri = publisher.publish(
+        ...     name="my_lens",
+        ...     source_schema_uri="at://did:plc:abc/ac.foundation.dataset.sampleSchema/source",
+        ...     target_schema_uri="at://did:plc:abc/ac.foundation.dataset.sampleSchema/target",
+        ...     code_repository="https://github.com/user/repo",
+        ...     code_commit="abc123def456",
+        ...     getter_path="mymodule.lenses:my_lens",
+        ...     putter_path="mymodule.lenses:my_lens_putter",
+        ... )
 
     Security Note:
         Lens code is stored as references to git repositories rather than
@@ -195,16 +194,14 @@ class LensLoader:
     using a lens requires installing the referenced code and importing
     it manually.
 
-    Example:
-        ::
-
-            >>> client = AtmosphereClient()
-            >>> loader = LensLoader(client)
-            >>>
-            >>> record = loader.get("at://did:plc:abc/ac.foundation.dataset.lens/xyz")
-            >>> print(record["name"])
-            >>> print(record["sourceSchema"])
-            >>> print(record.get("getterCode", {}).get("repository"))
+    Examples:
+        >>> client = AtmosphereClient()
+        >>> loader = LensLoader(client)
+        >>>
+        >>> record = loader.get("at://did:plc:abc/ac.foundation.dataset.lens/xyz")
+        >>> print(record["name"])
+        >>> print(record["sourceSchema"])
+        >>> print(record.get("getterCode", {}).get("repository"))
     """
 
     def __init__(self, client: AtmosphereClient):