atdata 0.1.3b4__py3-none-any.whl → 0.2.2b1__py3-none-any.whl

atdata/atmosphere/__init__.py

@@ -0,0 +1,332 @@
+"""ATProto integration for distributed dataset federation.
+
+This module provides ATProto publishing and discovery capabilities for atdata,
+enabling a loose federation of distributed, typed datasets on the AT Protocol
+network.
+
+Key components:
+
+- ``AtmosphereClient``: Authentication and session management for ATProto
+- ``SchemaPublisher``: Publish PackableSample schemas as ATProto records
+- ``DatasetPublisher``: Publish dataset index records with WebDataset URLs
+- ``LensPublisher``: Publish lens transformation records
+
+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")
+
+Note:
+    This module requires the ``atproto`` package to be installed::
+
+        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,
+    DatasetRecord,
+    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",
+    # Dataset operations
+    "DatasetPublisher",
+    "DatasetLoader",
+    # Lens operations
+    "LensPublisher",
+    "LensLoader",
+    # Types
+    "AtUri",
+    "SchemaRecord",
+    "DatasetRecord",
+    "LensRecord",
+]

atdata/atmosphere/_types.py

@@ -0,0 +1,331 @@
+"""Type definitions for ATProto record structures.
+
+This module defines the data structures used to represent ATProto records
+for schemas, datasets, and lenses. These types map to the Lexicon definitions
+in the ``ac.foundation.dataset.*`` namespace.
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Optional, Literal, Any
+
+# Lexicon namespace for atdata records
+LEXICON_NAMESPACE = "ac.foundation.dataset"
+
+
+@dataclass
+class AtUri:
+    """Parsed AT Protocol URI.
+
+    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'
+    """
+
+    authority: str
+    """The DID or handle of the repository owner."""
+
+    collection: str
+    """The NSID of the record collection."""
+
+    rkey: str
+    """The record key within the collection."""
+
+    @classmethod
+    def parse(cls, uri: str) -> "AtUri":
+        """Parse an AT URI string into components.
+
+        Args:
+            uri: AT URI string in format ``at://<authority>/<collection>/<rkey>``
+
+        Returns:
+            Parsed AtUri instance.
+
+        Raises:
+            ValueError: If the URI format is invalid.
+        """
+        if not uri.startswith("at://"):
+            raise ValueError(f"Invalid AT URI: must start with 'at://': {uri}")
+
+        parts = uri[5:].split("/")
+        if len(parts) < 3:
+            raise ValueError(f"Invalid AT URI: expected authority/collection/rkey: {uri}")
+
+        return cls(
+            authority=parts[0],
+            collection=parts[1],
+            rkey="/".join(parts[2:]),  # rkey may contain slashes
+        )
+
+    def __str__(self) -> str:
+        """Format as AT URI string."""
+        return f"at://{self.authority}/{self.collection}/{self.rkey}"
+
+
+@dataclass
+class FieldType:
+    """Schema field type definition.
+
+    Represents a type in the schema type system, supporting primitives,
+    ndarrays, and references to other schemas.
+    """
+
+    kind: Literal["primitive", "ndarray", "ref", "array"]
+    """The category of type."""
+
+    primitive: Optional[str] = None
+    """For kind='primitive': one of 'str', 'int', 'float', 'bool', 'bytes'."""
+
+    dtype: Optional[str] = None
+    """For kind='ndarray': numpy dtype string (e.g., 'float32')."""
+
+    shape: Optional[list[int | None]] = None
+    """For kind='ndarray': shape constraints (None for any dimension)."""
+
+    ref: Optional[str] = None
+    """For kind='ref': AT URI of referenced schema."""
+
+    items: Optional["FieldType"] = None
+    """For kind='array': type of array elements."""
+
+
+@dataclass
+class FieldDef:
+    """Schema field definition."""
+
+    name: str
+    """Field name."""
+
+    field_type: FieldType
+    """Type of this field."""
+
+    optional: bool = False
+    """Whether this field can be None."""
+
+    description: Optional[str] = None
+    """Human-readable description."""
+
+
+@dataclass
+class SchemaRecord:
+    """ATProto record for a PackableSample schema.
+
+    Maps to the ``ac.foundation.dataset.sampleSchema`` Lexicon.
+    """
+
+    name: str
+    """Human-readable schema name."""
+
+    version: str
+    """Semantic version string (e.g., '1.0.0')."""
+
+    fields: list[FieldDef]
+    """List of field definitions."""
+
+    description: Optional[str] = None
+    """Human-readable description."""
+
+    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    """When this record was created."""
+
+    metadata: Optional[dict] = None
+    """Arbitrary metadata as msgpack-encoded bytes."""
+
+    def to_record(self) -> dict:
+        """Convert to ATProto record dict for publishing."""
+        record = {
+            "$type": f"{LEXICON_NAMESPACE}.sampleSchema",
+            "name": self.name,
+            "version": self.version,
+            "fields": [self._field_to_dict(f) for f in self.fields],
+            "createdAt": self.created_at.isoformat(),
+        }
+        if self.description:
+            record["description"] = self.description
+        if self.metadata:
+            record["metadata"] = self.metadata
+        return record
+
+    def _field_to_dict(self, field_def: FieldDef) -> dict:
+        """Convert a field definition to dict."""
+        result = {
+            "name": field_def.name,
+            "fieldType": self._type_to_dict(field_def.field_type),
+            "optional": field_def.optional,
+        }
+        if field_def.description:
+            result["description"] = field_def.description
+        return result
+
+    def _type_to_dict(self, field_type: FieldType) -> dict:
+        """Convert a field type to dict."""
+        result: dict = {"$type": f"{LEXICON_NAMESPACE}.schemaType#{field_type.kind}"}
+
+        if field_type.kind == "primitive":
+            result["primitive"] = field_type.primitive
+        elif field_type.kind == "ndarray":
+            result["dtype"] = field_type.dtype
+            if field_type.shape:
+                result["shape"] = field_type.shape
+        elif field_type.kind == "ref":
+            result["ref"] = field_type.ref
+        elif field_type.kind == "array":
+            if field_type.items:
+                result["items"] = self._type_to_dict(field_type.items)
+
+        return result
+
+
+@dataclass
+class StorageLocation:
+    """Dataset storage location specification."""
+
+    kind: Literal["external", "blobs"]
+    """Storage type: external URLs or ATProto blobs."""
+
+    urls: Optional[list[str]] = None
+    """For kind='external': WebDataset URLs with brace notation."""
+
+    blob_refs: Optional[list[dict]] = None
+    """For kind='blobs': ATProto blob references."""
+
+
+@dataclass
+class DatasetRecord:
+    """ATProto record for a dataset index.
+
+    Maps to the ``ac.foundation.dataset.record`` Lexicon.
+    """
+
+    name: str
+    """Human-readable dataset name."""
+
+    schema_ref: str
+    """AT URI of the schema record."""
+
+    storage: StorageLocation
+    """Where the dataset data is stored."""
+
+    description: Optional[str] = None
+    """Human-readable description."""
+
+    tags: list[str] = field(default_factory=list)
+    """Searchable tags."""
+
+    license: Optional[str] = None
+    """SPDX license identifier."""
+
+    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    """When this record was created."""
+
+    metadata: Optional[bytes] = None
+    """Arbitrary metadata as msgpack-encoded bytes."""
+
+    def to_record(self) -> dict:
+        """Convert to ATProto record dict for publishing."""
+        record = {
+            "$type": f"{LEXICON_NAMESPACE}.record",
+            "name": self.name,
+            "schemaRef": self.schema_ref,
+            "storage": self._storage_to_dict(),
+            "createdAt": self.created_at.isoformat(),
+        }
+        if self.description:
+            record["description"] = self.description
+        if self.tags:
+            record["tags"] = self.tags
+        if self.license:
+            record["license"] = self.license
+        if self.metadata:
+            record["metadata"] = self.metadata
+        return record
+
+    def _storage_to_dict(self) -> dict:
+        """Convert storage location to dict."""
+        if self.storage.kind == "external":
+            return {
+                "$type": f"{LEXICON_NAMESPACE}.storageExternal",
+                "urls": self.storage.urls or [],
+            }
+        else:
+            return {
+                "$type": f"{LEXICON_NAMESPACE}.storageBlobs",
+                "blobs": self.storage.blob_refs or [],
+            }
+
+
+@dataclass
+class CodeReference:
+    """Reference to lens code in a git repository."""
+
+    repository: str
+    """Git repository URL."""
+
+    commit: str
+    """Git commit hash."""
+
+    path: str
+    """Path to the code file/function."""
+
+
+@dataclass
+class LensRecord:
+    """ATProto record for a lens transformation.
+
+    Maps to the ``ac.foundation.dataset.lens`` Lexicon.
+    """
+
+    name: str
+    """Human-readable lens name."""
+
+    source_schema: str
+    """AT URI of the source schema."""
+
+    target_schema: str
+    """AT URI of the target schema."""
+
+    description: Optional[str] = None
+    """What this transformation does."""
+
+    getter_code: Optional[CodeReference] = None
+    """Reference to getter function code."""
+
+    putter_code: Optional[CodeReference] = None
+    """Reference to putter function code."""
+
+    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    """When this record was created."""
+
+    def to_record(self) -> dict:
+        """Convert to ATProto record dict for publishing."""
+        record: dict[str, Any] = {
+            "$type": f"{LEXICON_NAMESPACE}.lens",
+            "name": self.name,
+            "sourceSchema": self.source_schema,
+            "targetSchema": self.target_schema,
+            "createdAt": self.created_at.isoformat(),
+        }
+        if self.description:
+            record["description"] = self.description
+        if self.getter_code:
+            record["getterCode"] = {
+                "repository": self.getter_code.repository,
+                "commit": self.getter_code.commit,
+                "path": self.getter_code.path,
+            }
+        if self.putter_code:
+            record["putterCode"] = {
+                "repository": self.putter_code.repository,
+                "commit": self.putter_code.commit,
+                "path": self.putter_code.path,
+            }
+        return record