atdata 0.2.0a1__py3-none-any.whl → 0.2.3b1__py3-none-any.whl

atdata/atmosphere/__init__.py

@@ -15,7 +15,7 @@ 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:
+Examples:
     >>> from atdata.atmosphere import AtmosphereClient, SchemaPublisher
     >>>
     >>> client = AtmosphereClient()
@@ -30,10 +30,13 @@ Note:
         pip install atproto
 """
 
+from typing import Iterator, Optional, Type, TYPE_CHECKING
+
 from .client import AtmosphereClient
 from .schema import SchemaPublisher, SchemaLoader
 from .records import DatasetPublisher, DatasetLoader
 from .lens import LensPublisher, LensLoader
+from .store import PDSBlobStore
 from ._types import (
     AtUri,
     SchemaRecord,
@@ -41,9 +44,274 @@ from ._types import (
     LensRecord,
 )
 
+if TYPE_CHECKING:
+    from ..dataset import Dataset
+    from .._protocols import Packable
+
+
+class AtmosphereIndexEntry:
+    """Entry wrapper for ATProto dataset records implementing IndexEntry protocol.
+
+    Attributes:
+        _uri: AT URI of the record.
+        _record: Raw record dictionary.
+    """
+
+    def __init__(self, uri: str, record: dict):
+        self._uri = uri
+        self._record = record
+
+    @property
+    def name(self) -> str:
+        """Human-readable dataset name."""
+        return self._record.get("name", "")
+
+    @property
+    def schema_ref(self) -> str:
+        """AT URI of the schema record."""
+        return self._record.get("schemaRef", "")
+
+    @property
+    def data_urls(self) -> list[str]:
+        """WebDataset URLs from external storage."""
+        storage = self._record.get("storage", {})
+        storage_type = storage.get("$type", "")
+        if "storageExternal" in storage_type:
+            return storage.get("urls", [])
+        return []
+
+    @property
+    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
+        return msgpack.unpackb(metadata_bytes, raw=False)
+
+    @property
+    def uri(self) -> str:
+        """AT URI of this record."""
+        return self._uri
+
+
+class AtmosphereIndex:
+    """ATProto index implementing AbstractIndex protocol.
+
+    Wraps SchemaPublisher/Loader and DatasetPublisher/Loader to provide
+    a unified interface compatible with LocalIndex.
+
+    Optionally accepts a ``PDSBlobStore`` for writing dataset shards as
+    ATProto blobs, enabling fully decentralized dataset storage.
+
+    Examples:
+        >>> 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")
+    """
+
+    def __init__(
+        self,
+        client: AtmosphereClient,
+        *,
+        data_store: Optional[PDSBlobStore] = None,
+    ):
+        """Initialize the atmosphere index.
+
+        Args:
+            client: Authenticated AtmosphereClient instance.
+            data_store: Optional PDSBlobStore for writing shards as blobs.
+                If provided, insert_dataset will upload shards to PDS.
+        """
+        self.client = client
+        self._schema_publisher = SchemaPublisher(client)
+        self._schema_loader = SchemaLoader(client)
+        self._dataset_publisher = DatasetPublisher(client)
+        self._dataset_loader = DatasetLoader(client)
+        self._data_store = data_store
+
+    @property
+    def data_store(self) -> Optional[PDSBlobStore]:
+        """The PDS blob store for writing shards, or None if not configured."""
+        return self._data_store
+
+    # Dataset operations
+
+    def insert_dataset(
+        self,
+        ds: "Dataset",
+        *,
+        name: str,
+        schema_ref: Optional[str] = None,
+        **kwargs,
+    ) -> AtmosphereIndexEntry:
+        """Insert a dataset into ATProto.
+
+        Args:
+            ds: The Dataset to publish.
+            name: Human-readable name.
+            schema_ref: Optional schema AT URI. If None, auto-publishes schema.
+            **kwargs: Additional options (description, tags, license).
+
+        Returns:
+            AtmosphereIndexEntry for the inserted dataset.
+        """
+        uri = self._dataset_publisher.publish(
+            ds,
+            name=name,
+            schema_uri=schema_ref,
+            description=kwargs.get("description"),
+            tags=kwargs.get("tags"),
+            license=kwargs.get("license"),
+            auto_publish_schema=(schema_ref is None),
+        )
+        record = self._dataset_loader.get(uri)
+        return AtmosphereIndexEntry(str(uri), record)
+
+    def get_dataset(self, ref: str) -> AtmosphereIndexEntry:
+        """Get a dataset by AT URI.
+
+        Args:
+            ref: AT URI of the dataset record.
+
+        Returns:
+            AtmosphereIndexEntry for the dataset.
+
+        Raises:
+            ValueError: If record is not a dataset.
+        """
+        record = self._dataset_loader.get(ref)
+        return AtmosphereIndexEntry(ref, record)
+
+    @property
+    def datasets(self) -> Iterator[AtmosphereIndexEntry]:
+        """Lazily iterate over all dataset entries (AbstractIndex protocol).
+
+        Uses the authenticated user's repository.
+
+        Yields:
+            AtmosphereIndexEntry for each dataset.
+        """
+        records = self._dataset_loader.list_all()
+        for rec in records:
+            uri = rec.get("uri", "")
+            yield AtmosphereIndexEntry(uri, rec.get("value", rec))
+
+    def list_datasets(self, repo: Optional[str] = None) -> list[AtmosphereIndexEntry]:
+        """Get all dataset entries as a materialized list (AbstractIndex protocol).
+
+        Args:
+            repo: DID of repository. Defaults to authenticated user.
+
+        Returns:
+            List of AtmosphereIndexEntry for each dataset.
+        """
+        records = self._dataset_loader.list_all(repo=repo)
+        return [
+            AtmosphereIndexEntry(rec.get("uri", ""), rec.get("value", rec))
+            for rec in records
+        ]
+
+    # Schema operations
+
+    def publish_schema(
+        self,
+        sample_type: "Type[Packable]",
+        *,
+        version: str = "1.0.0",
+        **kwargs,
+    ) -> str:
+        """Publish a schema to ATProto.
+
+        Args:
+            sample_type: A Packable type (PackableSample subclass or @packable-decorated).
+            version: Semantic version string.
+            **kwargs: Additional options (description, metadata).
+
+        Returns:
+            AT URI of the schema record.
+        """
+        uri = self._schema_publisher.publish(
+            sample_type,
+            version=version,
+            description=kwargs.get("description"),
+            metadata=kwargs.get("metadata"),
+        )
+        return str(uri)
+
+    def get_schema(self, ref: str) -> dict:
+        """Get a schema record by AT URI.
+
+        Args:
+            ref: AT URI of the schema record.
+
+        Returns:
+            Schema record dictionary.
+
+        Raises:
+            ValueError: If record is not a schema.
+        """
+        return self._schema_loader.get(ref)
+
+    @property
+    def schemas(self) -> Iterator[dict]:
+        """Lazily iterate over all schema records (AbstractIndex protocol).
+
+        Uses the authenticated user's repository.
+
+        Yields:
+            Schema records as dictionaries.
+        """
+        records = self._schema_loader.list_all()
+        for rec in records:
+            yield rec.get("value", rec)
+
+    def list_schemas(self, repo: Optional[str] = None) -> list[dict]:
+        """Get all schema records as a materialized list (AbstractIndex protocol).
+
+        Args:
+            repo: DID of repository. Defaults to authenticated user.
+
+        Returns:
+            List of schema records as dictionaries.
+        """
+        records = self._schema_loader.list_all(repo=repo)
+        return [rec.get("value", rec) for rec in records]
+
+    def decode_schema(self, ref: str) -> "Type[Packable]":
+        """Reconstruct a Python type from a schema record.
+
+        Args:
+            ref: AT URI of the schema record.
+
+        Returns:
+            Dynamically generated Packable type.
+
+        Raises:
+            ValueError: If schema cannot be decoded.
+        """
+        from .._schema_codec import schema_to_type
+
+        schema = self.get_schema(ref)
+        return schema_to_type(schema)
+
+
 __all__ = [
     # Client
     "AtmosphereClient",
+    # Storage
+    "PDSBlobStore",
+    # Unified index (AbstractIndex protocol)
+    "AtmosphereIndex",
+    "AtmosphereIndexEntry",
     # Schema operations
     "SchemaPublisher",
     "SchemaLoader",

atdata/atmosphere/_types.py

@@ -19,7 +19,7 @@ class AtUri:
 
     AT URIs follow the format: at://<authority>/<collection>/<rkey>
 
-    Example:
+    Examples:
         >>> uri = AtUri.parse("at://did:plc:abc123/ac.foundation.dataset.sampleSchema/xyz")
         >>> uri.authority
         'did:plc:abc123'
@@ -56,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,7 +34,7 @@ class AtmosphereClient:
     This class wraps the atproto SDK client and provides higher-level methods
     for working with atdata records (schemas, datasets, lenses).
 
-    Example:
+    Examples:
         >>> client = AtmosphereClient()
         >>> client.login("alice.bsky.social", "app-password")
         >>> print(client.did)
@@ -254,7 +255,18 @@ class AtmosphereClient:
             }
         )
 
-        return response.value
+        # Convert ATProto model to dict if needed
+        value = response.value
+        # DotDict and similar ATProto models have to_dict()
+        if hasattr(value, "to_dict") and callable(value.to_dict):
+            return value.to_dict()
+        elif isinstance(value, dict):
+            return dict(value)
+        elif hasattr(value, "model_dump") and callable(value.model_dump):
+            return value.model_dump()
+        elif hasattr(value, "__dict__"):
+            return dict(value.__dict__)
+        return value
 
     def delete_record(
         self,
@@ -287,6 +299,123 @@ class AtmosphereClient:
 
         self._client.com.atproto.repo.delete_record(data=data)
 
+    def upload_blob(
+        self,
+        data: bytes,
+        mime_type: str = "application/octet-stream",
+    ) -> dict:
+        """Upload binary data as a blob to the PDS.
+
+        Args:
+            data: Binary data to upload.
+            mime_type: MIME type of the data (for reference, not enforced by PDS).
+
+        Returns:
+            A blob reference dict with keys: '$type', 'ref', 'mimeType', 'size'.
+            This can be embedded directly in record fields.
+
+        Raises:
+            ValueError: If not authenticated.
+            atproto.exceptions.AtProtocolError: If upload fails.
+        """
+        self._ensure_authenticated()
+
+        response = self._client.upload_blob(data)
+        blob_ref = response.blob
+
+        # 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)
+            },
+            "mimeType": blob_ref.mime_type,
+            "size": blob_ref.size,
+        }
+
+    def get_blob(
+        self,
+        did: str,
+        cid: str,
+    ) -> bytes:
+        """Download a blob from a PDS.
+
+        This resolves the PDS endpoint from the DID document and fetches
+        the blob directly from the PDS.
+
+        Args:
+            did: The DID of the repository containing the blob.
+            cid: The CID of the blob.
+
+        Returns:
+            The blob data as bytes.
+
+        Raises:
+            ValueError: If PDS endpoint cannot be resolved.
+            requests.HTTPError: If blob fetch fails.
+        """
+        import requests
+
+        # Resolve PDS endpoint from DID document
+        pds_endpoint = self._resolve_pds_endpoint(did)
+        if not pds_endpoint:
+            raise ValueError(f"Could not resolve PDS endpoint for {did}")
+
+        # Fetch blob from PDS
+        url = f"{pds_endpoint}/xrpc/com.atproto.sync.getBlob"
+        response = requests.get(url, params={"did": did, "cid": cid})
+        response.raise_for_status()
+        return response.content
+
+    def _resolve_pds_endpoint(self, did: str) -> Optional[str]:
+        """Resolve the PDS endpoint for a DID.
+
+        Args:
+            did: The DID to resolve.
+
+        Returns:
+            The PDS service endpoint URL, or None if not found.
+        """
+        import requests
+
+        # For did:plc, query the PLC directory
+        if did.startswith("did:plc:"):
+            try:
+                response = requests.get(f"https://plc.directory/{did}")
+                response.raise_for_status()
+                did_doc = response.json()
+
+                for service in did_doc.get("service", []):
+                    if service.get("type") == "AtprotoPersonalDataServer":
+                        return service.get("serviceEndpoint")
+            except requests.RequestException:
+                return None
+
+        # For did:web, would need different resolution (not implemented)
+        return None
+
+    def get_blob_url(self, did: str, cid: str) -> str:
+        """Get the direct URL for fetching a blob.
+
+        This is useful for passing to WebDataset or other HTTP clients.
+
+        Args:
+            did: The DID of the repository containing the blob.
+            cid: The CID of the blob.
+
+        Returns:
+            The full URL for fetching the blob.
+
+        Raises:
+            ValueError: If PDS endpoint cannot be resolved.
+        """
+        pds_endpoint = self._resolve_pds_endpoint(did)
+        if not pds_endpoint:
+            raise ValueError(f"Could not resolve PDS endpoint for {did}")
+        return f"{pds_endpoint}/xrpc/com.atproto.sync.getBlob?did={did}&cid={cid}"
+
     def list_records(
         self,
         collection: str,
@@ -324,7 +453,21 @@ class AtmosphereClient:
             }
         )
 
-        records = [r.value for r in response.records]
+        # Convert ATProto models to dicts if needed
+        records = []
+        for r in response.records:
+            value = r.value
+            # DotDict and similar ATProto models have to_dict()
+            if hasattr(value, "to_dict") and callable(value.to_dict):
+                records.append(value.to_dict())
+            elif isinstance(value, dict):
+                records.append(dict(value))
+            elif hasattr(value, "model_dump") and callable(value.model_dump):
+                records.append(value.model_dump())
+            elif hasattr(value, "__dict__"):
+                records.append(dict(value.__dict__))
+            else:
+                records.append(value)
         return records, response.cursor
 
     # Convenience methods for atdata collections

atdata/atmosphere/lens.py

@@ -9,7 +9,7 @@ Note:
     implementations.
 """
 
-from typing import Optional, Callable
+from typing import Optional
 
 from .client import AtmosphereClient
 from ._types import (
@@ -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,7 +32,7 @@ class LensPublisher:
     This class creates lens records that reference source and target schemas
     and point to the transformation code in a git repository.
 
-    Example:
+    Examples:
         >>> @atdata.lens
         ... def my_lens(source: SourceType) -> TargetType:
         ...     return TargetType(field=source.other_field)
@@ -193,7 +194,7 @@ class LensLoader:
     using a lens requires installing the referenced code and importing
     it manually.
 
-    Example:
+    Examples:
         >>> client = AtmosphereClient()
         >>> loader = LensLoader(client)
         >>>