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

atdata/atmosphere/schema.py

@@ -0,0 +1,239 @@
+"""Schema publishing and loading for ATProto.
+
+This module provides classes for publishing PackableSample schemas to ATProto
+and loading them back. Schemas are published as ``ac.foundation.dataset.sampleSchema``
+records.
+"""
+
+from dataclasses import fields, is_dataclass
+from typing import Type, TypeVar, Optional, get_type_hints, get_origin, get_args
+
+from .client import AtmosphereClient
+from ._types import (
+    AtUri,
+    SchemaRecord,
+    FieldDef,
+    FieldType,
+    LEXICON_NAMESPACE,
+)
+from .._type_utils import (
+    numpy_dtype_to_string,
+    unwrap_optional,
+    is_ndarray_type,
+    extract_ndarray_dtype,
+)
+
+# Import for type checking only to avoid circular imports
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from ..dataset import PackableSample
+
+ST = TypeVar("ST", bound="PackableSample")
+
+
+class SchemaPublisher:
+    """Publishes PackableSample schemas to ATProto.
+
+    This class introspects a PackableSample class to extract its field
+    definitions and publishes them as an ATProto schema record.
+
+    Example:
+        ::
+
+            >>> @atdata.packable
+            ... class MySample:
+            ...     image: NDArray
+            ...     label: str
+            ...
+            >>> client = AtmosphereClient()
+            >>> client.login("handle", "password")
+            >>>
+            >>> publisher = SchemaPublisher(client)
+            >>> uri = publisher.publish(MySample, version="1.0.0")
+            >>> print(uri)
+            at://did:plc:.../ac.foundation.dataset.sampleSchema/...
+    """
+
+    def __init__(self, client: AtmosphereClient):
+        """Initialize the schema publisher.
+
+        Args:
+            client: Authenticated AtmosphereClient instance.
+        """
+        self.client = client
+
+    def publish(
+        self,
+        sample_type: Type[ST],
+        *,
+        name: Optional[str] = None,
+        version: str = "1.0.0",
+        description: Optional[str] = None,
+        metadata: Optional[dict] = None,
+        rkey: Optional[str] = None,
+    ) -> AtUri:
+        """Publish a PackableSample schema to ATProto.
+
+        Args:
+            sample_type: The PackableSample class to publish.
+            name: Human-readable name. Defaults to the class name.
+            version: Semantic version string (e.g., '1.0.0').
+            description: Human-readable description.
+            metadata: Arbitrary metadata dictionary.
+            rkey: Optional explicit record key. If not provided, a TID is generated.
+
+        Returns:
+            The AT URI of the created schema record.
+
+        Raises:
+            ValueError: If sample_type is not a dataclass or client is not authenticated.
+            TypeError: If a field type is not supported.
+        """
+        if not is_dataclass(sample_type):
+            raise ValueError(f"{sample_type.__name__} must be a dataclass (use @packable)")
+
+        # Build the schema record
+        schema_record = self._build_schema_record(
+            sample_type,
+            name=name,
+            version=version,
+            description=description,
+            metadata=metadata,
+        )
+
+        # Publish to ATProto
+        return self.client.create_record(
+            collection=f"{LEXICON_NAMESPACE}.sampleSchema",
+            record=schema_record.to_record(),
+            rkey=rkey,
+            validate=False,  # PDS doesn't know our lexicon
+        )
+
+    def _build_schema_record(
+        self,
+        sample_type: Type[ST],
+        *,
+        name: Optional[str],
+        version: str,
+        description: Optional[str],
+        metadata: Optional[dict],
+    ) -> SchemaRecord:
+        """Build a SchemaRecord from a PackableSample class."""
+        field_defs = []
+        type_hints = get_type_hints(sample_type)
+
+        for f in fields(sample_type):
+            field_type = type_hints.get(f.name, f.type)
+            field_def = self._field_to_def(f.name, field_type)
+            field_defs.append(field_def)
+
+        return SchemaRecord(
+            name=name or sample_type.__name__,
+            version=version,
+            description=description,
+            fields=field_defs,
+            metadata=metadata,
+        )
+
+    def _field_to_def(self, name: str, python_type) -> FieldDef:
+        """Convert a Python field to a FieldDef."""
+        python_type, is_optional = unwrap_optional(python_type)
+        field_type = self._python_type_to_field_type(python_type)
+        return FieldDef(name=name, field_type=field_type, optional=is_optional)
+
+    def _python_type_to_field_type(self, python_type) -> FieldType:
+        """Map a Python type to a FieldType."""
+        if python_type is str:
+            return FieldType(kind="primitive", primitive="str")
+        if python_type is int:
+            return FieldType(kind="primitive", primitive="int")
+        if python_type is float:
+            return FieldType(kind="primitive", primitive="float")
+        if python_type is bool:
+            return FieldType(kind="primitive", primitive="bool")
+        if python_type is bytes:
+            return FieldType(kind="primitive", primitive="bytes")
+
+        if is_ndarray_type(python_type):
+            return FieldType(kind="ndarray", dtype=extract_ndarray_dtype(python_type), shape=None)
+
+        origin = get_origin(python_type)
+        if origin is list:
+            args = get_args(python_type)
+            items = self._python_type_to_field_type(args[0]) if args else FieldType(kind="primitive", primitive="str")
+            return FieldType(kind="array", items=items)
+
+        if is_dataclass(python_type):
+            raise TypeError(
+                f"Nested dataclass types not yet supported: {python_type.__name__}. "
+                "Publish nested types separately and use references."
+            )
+
+        raise TypeError(f"Unsupported type for schema field: {python_type}")
+
+
+class SchemaLoader:
+    """Loads PackableSample schemas from ATProto.
+
+    This class fetches schema records from ATProto and can list available
+    schemas from a repository.
+
+    Example:
+        ::
+
+            >>> client = AtmosphereClient()
+            >>> client.login("handle", "password")
+            >>>
+            >>> loader = SchemaLoader(client)
+            >>> schema = loader.get("at://did:plc:.../ac.foundation.dataset.sampleSchema/...")
+            >>> print(schema["name"])
+            'MySample'
+    """
+
+    def __init__(self, client: AtmosphereClient):
+        """Initialize the schema loader.
+
+        Args:
+            client: AtmosphereClient instance (authentication optional for reads).
+        """
+        self.client = client
+
+    def get(self, uri: str | AtUri) -> dict:
+        """Fetch a schema record by AT URI.
+
+        Args:
+            uri: The AT URI of the schema record.
+
+        Returns:
+            The schema record as a dictionary.
+
+        Raises:
+            ValueError: If the record is not a schema record.
+            atproto.exceptions.AtProtocolError: If record not found.
+        """
+        record = self.client.get_record(uri)
+
+        expected_type = f"{LEXICON_NAMESPACE}.sampleSchema"
+        if record.get("$type") != expected_type:
+            raise ValueError(
+                f"Record at {uri} is not a schema record. "
+                f"Expected $type='{expected_type}', got '{record.get('$type')}'"
+            )
+
+        return record
+
+    def list_all(
+        self,
+        repo: Optional[str] = None,
+        limit: int = 100,
+    ) -> list[dict]:
+        """List schema records from a repository.
+
+        Args:
+            repo: The DID of the repository. Defaults to authenticated user.
+            limit: Maximum number of records to return.
+
+        Returns:
+            List of schema records.
+        """
+        return self.client.list_schemas(repo=repo, limit=limit)

atdata/atmosphere/store.py

@@ -0,0 +1,208 @@
+"""PDS blob storage for dataset shards.
+
+This module provides ``PDSBlobStore``, an implementation of the AbstractDataStore
+protocol that stores dataset shards as ATProto blobs in a Personal Data Server.
+
+This enables fully decentralized dataset storage where both metadata (records)
+and data (blobs) live on the AT Protocol network.
+
+Example:
+    ::
+
+        >>> from atdata.atmosphere import AtmosphereClient, PDSBlobStore
+        >>>
+        >>> client = AtmosphereClient()
+        >>> client.login("handle.bsky.social", "app-password")
+        >>>
+        >>> store = PDSBlobStore(client)
+        >>> urls = store.write_shards(dataset, prefix="mnist/v1")
+        >>> print(urls)
+        ['at://did:plc:.../blob/bafyrei...', ...]
+"""
+
+from __future__ import annotations
+
+import io
+import tempfile
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import webdataset as wds
+
+if TYPE_CHECKING:
+    from ..dataset import Dataset
+    from .client import AtmosphereClient
+
+
+@dataclass
+class PDSBlobStore:
+    """PDS blob store implementing AbstractDataStore protocol.
+
+    Stores dataset shards as ATProto blobs, enabling decentralized dataset
+    storage on the AT Protocol network.
+
+    Each shard is written to a temporary tar file, then uploaded as a blob
+    to the user's PDS. The returned URLs are AT URIs that can be resolved
+    to HTTP URLs for streaming.
+
+    Attributes:
+        client: Authenticated AtmosphereClient instance.
+
+    Example:
+        ::
+
+            >>> store = PDSBlobStore(client)
+            >>> urls = store.write_shards(dataset, prefix="training/v1")
+            >>> # Returns AT URIs like:
+            >>> # ['at://did:plc:abc/blob/bafyrei...', ...]
+    """
+
+    client: "AtmosphereClient"
+
+    def write_shards(
+        self,
+        ds: "Dataset",
+        *,
+        prefix: str,
+        maxcount: int = 10000,
+        maxsize: float = 3e9,
+        **kwargs: Any,
+    ) -> list[str]:
+        """Write dataset shards as PDS blobs.
+
+        Creates tar archives from the dataset and uploads each as a blob
+        to the authenticated user's PDS.
+
+        Args:
+            ds: The Dataset to write.
+            prefix: Logical path prefix for naming (used in shard names only).
+            maxcount: Maximum samples per shard (default: 10000).
+            maxsize: Maximum shard size in bytes (default: 3GB, PDS limit).
+            **kwargs: Additional args passed to wds.ShardWriter.
+
+        Returns:
+            List of AT URIs for the written blobs, in format:
+            ``at://{did}/blob/{cid}``
+
+        Raises:
+            ValueError: If not authenticated.
+            RuntimeError: If no shards were written.
+
+        Note:
+            PDS blobs have size limits (typically 50MB-5GB depending on PDS).
+            Adjust maxcount/maxsize to stay within limits.
+        """
+        if not self.client.did:
+            raise ValueError("Client must be authenticated to upload blobs")
+
+        did = self.client.did
+        blob_urls: list[str] = []
+
+        # Write shards to temp files, upload each as blob
+        with tempfile.TemporaryDirectory() as temp_dir:
+            shard_pattern = f"{temp_dir}/shard-%06d.tar"
+            written_files: list[str] = []
+
+            # Track written files via custom post callback
+            def track_file(fname: str) -> None:
+                written_files.append(fname)
+
+            with wds.writer.ShardWriter(
+                shard_pattern,
+                maxcount=maxcount,
+                maxsize=maxsize,
+                post=track_file,
+                **kwargs,
+            ) as sink:
+                for sample in ds.ordered(batch_size=None):
+                    sink.write(sample.as_wds)
+
+            if not written_files:
+                raise RuntimeError("No shards written")
+
+            # Upload each shard as a blob
+            for shard_path in written_files:
+                with open(shard_path, "rb") as f:
+                    shard_data = f.read()
+
+                blob_ref = self.client.upload_blob(
+                    shard_data,
+                    mime_type="application/x-tar",
+                )
+
+                # Extract CID from blob reference
+                cid = blob_ref["ref"]["$link"]
+                at_uri = f"at://{did}/blob/{cid}"
+                blob_urls.append(at_uri)
+
+        return blob_urls
+
+    def read_url(self, url: str) -> str:
+        """Resolve an AT URI blob reference to an HTTP URL.
+
+        Transforms ``at://did/blob/cid`` URIs to HTTP URLs that can be
+        streamed by WebDataset.
+
+        Args:
+            url: AT URI in format ``at://{did}/blob/{cid}``.
+
+        Returns:
+            HTTP URL for fetching the blob via PDS API.
+
+        Raises:
+            ValueError: If URL format is invalid or PDS cannot be resolved.
+        """
+        if not url.startswith("at://"):
+            # Not an AT URI, return unchanged
+            return url
+
+        # Parse at://did/blob/cid
+        parts = url[5:].split("/")  # Remove 'at://'
+        if len(parts) != 3 or parts[1] != "blob":
+            raise ValueError(f"Invalid blob AT URI format: {url}")
+
+        did, _, cid = parts
+        return self.client.get_blob_url(did, cid)
+
+    def supports_streaming(self) -> bool:
+        """PDS blobs support streaming via HTTP.
+
+        Returns:
+            True.
+        """
+        return True
+
+    def create_source(self, urls: list[str]) -> "BlobSource":
+        """Create a BlobSource for reading these AT URIs.
+
+        This is a convenience method for creating a DataSource that can
+        stream the blobs written by this store.
+
+        Args:
+            urls: List of AT URIs from write_shards().
+
+        Returns:
+            BlobSource configured for the given URLs.
+
+        Raises:
+            ValueError: If URLs are not valid AT URIs.
+        """
+        from .._sources import BlobSource
+
+        blob_refs: list[dict[str, str]] = []
+
+        for url in urls:
+            if not url.startswith("at://"):
+                raise ValueError(f"Not an AT URI: {url}")
+
+            parts = url[5:].split("/")
+            if len(parts) != 3 or parts[1] != "blob":
+                raise ValueError(f"Invalid blob AT URI: {url}")
+
+            did, _, cid = parts
+            blob_refs.append({"did": did, "cid": cid})
+
+        return BlobSource(blob_refs=blob_refs)
+
+
+__all__ = ["PDSBlobStore"]

atdata/cli/__init__.py

@@ -0,0 +1,213 @@
+"""Command-line interface for atdata.
+
+This module provides CLI commands for managing local development infrastructure
+and diagnosing configuration issues.
+
+Commands:
+    atdata local up     Start Redis and MinIO containers for local development
+    atdata local down   Stop local development containers
+    atdata diagnose     Check Redis configuration and connectivity
+    atdata version      Show version information
+
+Example:
+    $ atdata local up
+    Starting Redis on port 6379...
+    Starting MinIO on port 9000...
+    Local infrastructure ready.
+
+    $ atdata diagnose
+    Checking Redis configuration...
+    ✓ Redis connected
+    ✓ Persistence enabled (AOF)
+    ✓ Memory policy: noeviction
+"""
+
+import argparse
+import sys
+from typing import Sequence
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    """Main entry point for the atdata CLI.
+
+    Args:
+        argv: Command-line arguments. If None, uses sys.argv[1:].
+
+    Returns:
+        Exit code (0 for success, non-zero for errors).
+    """
+    parser = argparse.ArgumentParser(
+        prog="atdata",
+        description="A loose federation of distributed, typed datasets",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        "--version", "-v",
+        action="store_true",
+        help="Show version information",
+    )
+
+    subparsers = parser.add_subparsers(dest="command", help="Available commands")
+
+    # 'local' command group
+    local_parser = subparsers.add_parser(
+        "local",
+        help="Manage local development infrastructure",
+    )
+    local_subparsers = local_parser.add_subparsers(
+        dest="local_command",
+        help="Local infrastructure commands",
+    )
+
+    # 'local up' command
+    up_parser = local_subparsers.add_parser(
+        "up",
+        help="Start Redis and MinIO containers",
+    )
+    up_parser.add_argument(
+        "--redis-port",
+        type=int,
+        default=6379,
+        help="Redis port (default: 6379)",
+    )
+    up_parser.add_argument(
+        "--minio-port",
+        type=int,
+        default=9000,
+        help="MinIO API port (default: 9000)",
+    )
+    up_parser.add_argument(
+        "--minio-console-port",
+        type=int,
+        default=9001,
+        help="MinIO console port (default: 9001)",
+    )
+    up_parser.add_argument(
+        "--detach", "-d",
+        action="store_true",
+        default=True,
+        help="Run containers in detached mode (default: True)",
+    )
+
+    # 'local down' command
+    down_parser = local_subparsers.add_parser(
+        "down",
+        help="Stop local development containers",
+    )
+    down_parser.add_argument(
+        "--volumes", "-v",
+        action="store_true",
+        help="Also remove volumes (deletes all data)",
+    )
+
+    # 'local status' command
+    local_subparsers.add_parser(
+        "status",
+        help="Show status of local infrastructure",
+    )
+
+    # 'diagnose' command
+    diagnose_parser = subparsers.add_parser(
+        "diagnose",
+        help="Diagnose Redis configuration and connectivity",
+    )
+    diagnose_parser.add_argument(
+        "--host",
+        default="localhost",
+        help="Redis host (default: localhost)",
+    )
+    diagnose_parser.add_argument(
+        "--port",
+        type=int,
+        default=6379,
+        help="Redis port (default: 6379)",
+    )
+
+    # 'version' command (alternative to --version flag)
+    subparsers.add_parser(
+        "version",
+        help="Show version information",
+    )
+
+    args = parser.parse_args(argv)
+
+    # Handle --version flag
+    if args.version or args.command == "version":
+        return _cmd_version()
+
+    # Handle 'local' commands
+    if args.command == "local":
+        if args.local_command == "up":
+            return _cmd_local_up(
+                redis_port=args.redis_port,
+                minio_port=args.minio_port,
+                minio_console_port=args.minio_console_port,
+                detach=args.detach,
+            )
+        elif args.local_command == "down":
+            return _cmd_local_down(remove_volumes=args.volumes)
+        elif args.local_command == "status":
+            return _cmd_local_status()
+        else:
+            local_parser.print_help()
+            return 1
+
+    # Handle 'diagnose' command
+    if args.command == "diagnose":
+        return _cmd_diagnose(host=args.host, port=args.port)
+
+    # No command given
+    parser.print_help()
+    return 0
+
+
+def _cmd_version() -> int:
+    """Show version information."""
+    try:
+        from atdata import __version__
+        version = __version__
+    except ImportError:
+        # Fallback to package metadata
+        from importlib.metadata import version as pkg_version
+        version = pkg_version("atdata")
+
+    print(f"atdata {version}")
+    return 0
+
+
+def _cmd_local_up(
+    redis_port: int,
+    minio_port: int,
+    minio_console_port: int,
+    detach: bool,
+) -> int:
+    """Start local development infrastructure."""
+    from .local import local_up
+    return local_up(
+        redis_port=redis_port,
+        minio_port=minio_port,
+        minio_console_port=minio_console_port,
+        detach=detach,
+    )
+
+
+def _cmd_local_down(remove_volumes: bool) -> int:
+    """Stop local development infrastructure."""
+    from .local import local_down
+    return local_down(remove_volumes=remove_volumes)
+
+
+def _cmd_local_status() -> int:
+    """Show status of local infrastructure."""
+    from .local import local_status
+    return local_status()
+
+
+def _cmd_diagnose(host: str, port: int) -> int:
+    """Diagnose Redis configuration."""
+    from .diagnose import diagnose_redis
+    return diagnose_redis(host=host, port=port)
+
+
+if __name__ == "__main__":
+    sys.exit(main())