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

atdata/atmosphere/__init__.py

@@ -16,13 +16,15 @@ 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")
+    ::
+
+        >>> 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::
@@ -30,10 +32,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 +46,275 @@ 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.
+
+    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")
+    """
+
+    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

@@ -20,13 +20,15 @@ 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'
+        ::
+
+            >>> 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

atdata/atmosphere/client.py

@@ -34,10 +34,12 @@ class AtmosphereClient:
     for working with atdata records (schemas, datasets, lenses).
 
     Example:
-        >>> client = AtmosphereClient()
-        >>> client.login("alice.bsky.social", "app-password")
-        >>> print(client.did)
-        'did:plc:...'
+        ::
+
+            >>> 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
@@ -254,7 +256,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 +300,119 @@ 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 +450,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 (
@@ -32,23 +32,25 @@ class LensPublisher:
     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",
-        ... )
+        ::
+
+            >>> @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
@@ -194,13 +196,15 @@ class LensLoader:
     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"))
+        ::
+
+            >>> 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):