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

atdata/repository.py

@@ -5,7 +5,7 @@ optional ``AbstractDataStore`` (shard storage), forming a named storage unit
 that can be mounted into an ``Index``.
 
 The ``_AtmosphereBackend`` is an internal adapter that wraps an
-``AtmosphereClient`` to present the same operational surface as a repository,
+``Atmosphere`` to present the same operational surface as a repository,
 but routes through the ATProto network instead of a local provider.
 
 Examples:
@@ -105,7 +105,7 @@ def create_repository(
 
 
 class _AtmosphereBackend:
-    """Internal adapter wrapping AtmosphereClient for Index routing.
+    """Internal adapter wrapping Atmosphere for Index routing.
 
     This class extracts the operational logic from ``AtmosphereIndex`` into an
     internal component that the unified ``Index`` uses for ATProto resolution.
@@ -117,15 +117,15 @@ class _AtmosphereBackend:
 
     def __init__(
         self,
-        client: Any,  # AtmosphereClient, typed as Any to avoid hard import
+        client: Any,  # Atmosphere, typed as Any to avoid hard import
         *,
         data_store: Optional[AbstractDataStore] = None,
     ) -> None:
-        from .atmosphere.client import AtmosphereClient
+        from .atmosphere.client import Atmosphere
 
-        if not isinstance(client, AtmosphereClient):
-            raise TypeError(f"Expected AtmosphereClient, got {type(client).__name__}")
-        self.client: AtmosphereClient = client
+        if not isinstance(client, Atmosphere):
+            raise TypeError(f"Expected Atmosphere, got {type(client).__name__}")
+        self.client: Atmosphere = client
         self._data_store = data_store
         self._schema_publisher: Any = None
         self._schema_loader: Any = None
@@ -210,14 +210,26 @@ class _AtmosphereBackend:
         *,
         name: str,
         schema_ref: str | None = None,
+        data_urls: list[str] | None = None,
+        blob_refs: list[dict] | None = None,
         **kwargs: Any,
     ) -> Any:
         """Insert a dataset into ATProto.
 
+        When *blob_refs* is provided the record uses ``storageBlobs`` with
+        embedded blob reference objects so the PDS retains the uploaded blobs.
+
+        When *data_urls* is provided (without *blob_refs*) the record uses
+        ``storageExternal`` with those URLs.
+
         Args:
             ds: The Dataset to publish.
             name: Human-readable name.
             schema_ref: Optional schema AT URI. If None, auto-publishes schema.
+            data_urls: Explicit shard URLs to store in the record.  When
+                provided, these replace whatever ``ds.url`` contains.
+            blob_refs: Pre-uploaded blob reference dicts from
+                ``PDSBlobStore``.  Takes precedence over *data_urls*.
             **kwargs: Additional options (description, tags, license).
 
         Returns:
@@ -226,15 +238,53 @@ class _AtmosphereBackend:
         self._ensure_loaders()
         from .atmosphere import AtmosphereIndexEntry
 
-        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),
-        )
+        if blob_refs is not None or data_urls is not None:
+            # Ensure schema is published first
+            if schema_ref is None:
+                from .atmosphere import SchemaPublisher
+
+                sp = SchemaPublisher(self.client)
+                schema_uri_obj = sp.publish(
+                    ds.sample_type,
+                    version=kwargs.get("schema_version", "1.0.0"),
+                )
+                schema_ref = str(schema_uri_obj)
+
+            metadata = kwargs.get("metadata")
+            if metadata is None and hasattr(ds, "_metadata"):
+                metadata = ds._metadata
+
+            if blob_refs is not None:
+                uri = self._dataset_publisher.publish_with_blob_refs(
+                    blob_refs=blob_refs,
+                    schema_uri=schema_ref,
+                    name=name,
+                    description=kwargs.get("description"),
+                    tags=kwargs.get("tags"),
+                    license=kwargs.get("license"),
+                    metadata=metadata,
+                )
+            else:
+                uri = self._dataset_publisher.publish_with_urls(
+                    urls=data_urls,
+                    schema_uri=schema_ref,
+                    name=name,
+                    description=kwargs.get("description"),
+                    tags=kwargs.get("tags"),
+                    license=kwargs.get("license"),
+                    metadata=metadata,
+                )
+        else:
+            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)
 

atdata/stores/__init__.py

@@ -0,0 +1,23 @@
+"""Data stores for atdata datasets.
+
+Key classes:
+
+- ``LocalDiskStore``: Local filesystem data store.
+- ``S3DataStore``: S3-compatible object storage.
+"""
+
+from atdata.stores._disk import LocalDiskStore
+from atdata.stores._s3 import (
+    S3DataStore,
+    _s3_env,
+    _s3_from_credentials,
+    _create_s3_write_callbacks,
+)
+
+__all__ = [
+    "LocalDiskStore",
+    "S3DataStore",
+    "_s3_env",
+    "_s3_from_credentials",
+    "_create_s3_write_callbacks",
+]

atdata/stores/_disk.py

@@ -0,0 +1,131 @@
+"""Local filesystem data store for WebDataset shards.
+
+Writes and reads WebDataset tar archives on the local filesystem,
+implementing the ``AbstractDataStore`` protocol.
+
+Examples:
+    >>> store = LocalDiskStore(root="~/.atdata/data")
+    >>> urls = store.write_shards(dataset, prefix="mnist/v1")
+    >>> print(urls[0])
+    /home/user/.atdata/data/mnist/v1/data--a1b2c3--000000.tar
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+from uuid import uuid4
+
+import webdataset as wds
+
+if TYPE_CHECKING:
+    from atdata.dataset import Dataset
+
+
+class LocalDiskStore:
+    """Local filesystem data store.
+
+    Writes WebDataset shards to a directory on disk. Implements the
+    ``AbstractDataStore`` protocol for use with ``Index``.
+
+    Args:
+        root: Root directory for shard storage. Defaults to
+            ``~/.atdata/data/``. Created automatically if it does
+            not exist.
+
+    Examples:
+        >>> store = LocalDiskStore()
+        >>> urls = store.write_shards(dataset, prefix="my-dataset")
+    """
+
+    def __init__(self, root: str | Path | None = None) -> None:
+        if root is None:
+            root = Path.home() / ".atdata" / "data"
+        self._root = Path(root).expanduser().resolve()
+        self._root.mkdir(parents=True, exist_ok=True)
+
+    @property
+    def root(self) -> Path:
+        """Root directory for shard storage."""
+        return self._root
+
+    def write_shards(
+        self,
+        ds: "Dataset",
+        *,
+        prefix: str,
+        **kwargs: Any,
+    ) -> list[str]:
+        """Write dataset shards to the local filesystem.
+
+        Args:
+            ds: The Dataset to write.
+            prefix: Path prefix within root (e.g., ``'datasets/mnist/v1'``).
+            **kwargs: Additional args passed to ``wds.writer.ShardWriter``
+                (e.g., ``maxcount``, ``maxsize``).
+
+        Returns:
+            List of absolute file paths for the written shards.
+
+        Raises:
+            RuntimeError: If no shards were written.
+        """
+        from atdata._logging import get_logger, log_operation
+
+        log = get_logger()
+        shard_dir = self._root / prefix
+        shard_dir.mkdir(parents=True, exist_ok=True)
+
+        new_uuid = str(uuid4())[:8]
+        shard_pattern = str(shard_dir / f"data--{new_uuid}--%06d.tar")
+
+        written_shards: list[str] = []
+
+        def _track_shard(path: str) -> None:
+            written_shards.append(str(Path(path).resolve()))
+
+        # Filter out kwargs that are specific to other stores (e.g. S3)
+        # and not understood by wds.writer.ShardWriter / TarWriter.
+        writer_kwargs = {k: v for k, v in kwargs.items() if k not in ("cache_local",)}
+
+        with log_operation("LocalDiskStore.write_shards", prefix=prefix):
+            with wds.writer.ShardWriter(
+                shard_pattern,
+                post=_track_shard,
+                **writer_kwargs,
+            ) as sink:
+                for sample in ds.ordered(batch_size=None):
+                    sink.write(sample.as_wds)
+
+            if not written_shards:
+                raise RuntimeError(
+                    f"No shards written for prefix {prefix!r} in {self._root}"
+                )
+
+            log.info(
+                "LocalDiskStore.write_shards: wrote %d shard(s)", len(written_shards)
+            )
+
+        return written_shards
+
+    def read_url(self, url: str) -> str:
+        """Resolve a storage URL for reading.
+
+        Local filesystem paths are returned as-is since WebDataset
+        can read them directly.
+
+        Args:
+            url: Absolute file path to a shard.
+
+        Returns:
+            The same path, unchanged.
+        """
+        return url
+
+    def supports_streaming(self) -> bool:
+        """Whether this store supports streaming reads.
+
+        Returns:
+            ``True`` — local filesystem supports streaming.
+        """
+        return True

atdata/{local → stores}/_s3.py

@@ -189,129 +189,151 @@ class S3DataStore:
         Raises:
             RuntimeError: If no shards were written.
         """
+        from atdata._logging import get_logger, log_operation
+
+        log = get_logger()
+
         new_uuid = str(uuid4())
         shard_pattern = f"{self.bucket}/{prefix}/data--{new_uuid}--%06d.tar"
 
         written_shards: list[str] = []
 
-        # Manifest tracking state shared with the post callback
-        manifest_builders: list = []
-        current_builder: list = [None]  # mutable ref for closure
-        shard_counter: list[int] = [0]
-
-        if manifest:
-            from atdata.manifest import ManifestBuilder, ManifestWriter
-
-            def _make_builder(shard_idx: int) -> ManifestBuilder:
-                shard_id = f"{self.bucket}/{prefix}/data--{new_uuid}--{shard_idx:06d}"
-                return ManifestBuilder(
-                    sample_type=ds.sample_type,
-                    shard_id=shard_id,
-                    schema_version=schema_version,
-                    source_job_id=source_job_id,
-                    parent_shards=parent_shards,
-                    pipeline_version=pipeline_version,
-                )
-
-            current_builder[0] = _make_builder(0)
-
-        with TemporaryDirectory() as temp_dir:
-            writer_opener, writer_post_orig = _create_s3_write_callbacks(
-                credentials=self.credentials,
-                temp_dir=temp_dir,
-                written_shards=written_shards,
-                fs=self._fs,
-                cache_local=cache_local,
-                add_s3_prefix=True,
-            )
+        with log_operation(
+            "S3DataStore.write_shards",
+            prefix=prefix,
+            bucket=self.bucket,
+            manifest=manifest,
+        ):
+            # Manifest tracking state shared with the post callback
+            manifest_builders: list = []
+            current_builder: list = [None]  # mutable ref for closure
+            shard_counter: list[int] = [0]
 
             if manifest:
+                from atdata.manifest import ManifestBuilder, ManifestWriter
+
+                def _make_builder(shard_idx: int) -> ManifestBuilder:
+                    shard_id = (
+                        f"{self.bucket}/{prefix}/data--{new_uuid}--{shard_idx:06d}"
+                    )
+                    return ManifestBuilder(
+                        sample_type=ds.sample_type,
+                        shard_id=shard_id,
+                        schema_version=schema_version,
+                        source_job_id=source_job_id,
+                        parent_shards=parent_shards,
+                        pipeline_version=pipeline_version,
+                    )
+
+                current_builder[0] = _make_builder(0)
+
+            with TemporaryDirectory() as temp_dir:
+                writer_opener, writer_post_orig = _create_s3_write_callbacks(
+                    credentials=self.credentials,
+                    temp_dir=temp_dir,
+                    written_shards=written_shards,
+                    fs=self._fs,
+                    cache_local=cache_local,
+                    add_s3_prefix=True,
+                )
 
-                def writer_post(p: str):
-                    # Finalize the current manifest builder when a shard completes
-                    builder = current_builder[0]
-                    if builder is not None:
-                        manifest_builders.append(builder)
-                    # Advance to the next shard's builder
-                    shard_counter[0] += 1
-                    current_builder[0] = _make_builder(shard_counter[0])
-                    # Call original post callback
-                    writer_post_orig(p)
-            else:
-                writer_post = writer_post_orig
-
-            offset = 0
-            with wds.writer.ShardWriter(
-                shard_pattern,
-                opener=writer_opener,
-                post=writer_post,
-                **kwargs,
-            ) as sink:
-                for sample in ds.ordered(batch_size=None):
-                    wds_dict = sample.as_wds
-                    sink.write(wds_dict)
-
-                    if manifest and current_builder[0] is not None:
-                        packed_size = len(wds_dict.get("msgpack", b""))
-                        current_builder[0].add_sample(
-                            key=wds_dict["__key__"],
-                            offset=offset,
-                            size=packed_size,
-                            sample=sample,
-                        )
-                        # Approximate tar entry: 512-byte header + data rounded to 512
-                        offset += 512 + packed_size + (512 - packed_size % 512) % 512
-
-            # Finalize the last shard's builder (post isn't called for the last shard
-            # until ShardWriter closes, but we handle it here for safety)
-            if manifest and current_builder[0] is not None:
-                builder = current_builder[0]
-                if builder._rows:  # Only if samples were added
-                    manifest_builders.append(builder)
-
-            # Write all manifest files
-            if manifest:
-                for builder in manifest_builders:
-                    built = builder.build()
-                    writer = ManifestWriter(Path(temp_dir) / Path(built.shard_id))
-                    json_path, parquet_path = writer.write(built)
-
-                    # Upload manifest files to S3 alongside shards
-                    shard_id = built.shard_id
-                    json_key = f"{shard_id}.manifest.json"
-                    parquet_key = f"{shard_id}.manifest.parquet"
-
-                    if cache_local:
-                        import boto3
-
-                        s3_kwargs = {
-                            "aws_access_key_id": self.credentials["AWS_ACCESS_KEY_ID"],
-                            "aws_secret_access_key": self.credentials[
-                                "AWS_SECRET_ACCESS_KEY"
-                            ],
-                        }
-                        if "AWS_ENDPOINT" in self.credentials:
-                            s3_kwargs["endpoint_url"] = self.credentials["AWS_ENDPOINT"]
-                        s3_client = boto3.client("s3", **s3_kwargs)
-
-                        bucket_name = Path(shard_id).parts[0]
-                        json_s3_key = str(Path(*Path(json_key).parts[1:]))
-                        parquet_s3_key = str(Path(*Path(parquet_key).parts[1:]))
-
-                        with open(json_path, "rb") as f:
-                            s3_client.put_object(
-                                Bucket=bucket_name, Key=json_s3_key, Body=f.read()
+                if manifest:
+
+                    def writer_post(p: str):
+                        # Finalize the current manifest builder when a shard completes
+                        builder = current_builder[0]
+                        if builder is not None:
+                            manifest_builders.append(builder)
+                        # Advance to the next shard's builder
+                        shard_counter[0] += 1
+                        current_builder[0] = _make_builder(shard_counter[0])
+                        # Call original post callback
+                        writer_post_orig(p)
+                else:
+                    writer_post = writer_post_orig
+
+                offset = 0
+                with wds.writer.ShardWriter(
+                    shard_pattern,
+                    opener=writer_opener,
+                    post=writer_post,
+                    **kwargs,
+                ) as sink:
+                    for sample in ds.ordered(batch_size=None):
+                        wds_dict = sample.as_wds
+                        sink.write(wds_dict)
+
+                        if manifest and current_builder[0] is not None:
+                            packed_size = len(wds_dict.get("msgpack", b""))
+                            current_builder[0].add_sample(
+                                key=wds_dict["__key__"],
+                                offset=offset,
+                                size=packed_size,
+                                sample=sample,
                             )
-                        with open(parquet_path, "rb") as f:
-                            s3_client.put_object(
-                                Bucket=bucket_name, Key=parquet_s3_key, Body=f.read()
+                            # Approximate tar entry: 512-byte header + data rounded to 512
+                            offset += (
+                                512 + packed_size + (512 - packed_size % 512) % 512
                             )
-                    else:
-                        self._fs.put(str(json_path), f"s3://{json_key}")
-                        self._fs.put(str(parquet_path), f"s3://{parquet_key}")
 
-        if len(written_shards) == 0:
-            raise RuntimeError("No shards written")
+                # Finalize the last shard's builder (post isn't called for the last shard
+                # until ShardWriter closes, but we handle it here for safety)
+                if manifest and current_builder[0] is not None:
+                    builder = current_builder[0]
+                    if builder._rows:  # Only if samples were added
+                        manifest_builders.append(builder)
+
+                # Write all manifest files
+                if manifest:
+                    for builder in manifest_builders:
+                        built = builder.build()
+                        writer = ManifestWriter(Path(temp_dir) / Path(built.shard_id))
+                        json_path, parquet_path = writer.write(built)
+
+                        # Upload manifest files to S3 alongside shards
+                        shard_id = built.shard_id
+                        json_key = f"{shard_id}.manifest.json"
+                        parquet_key = f"{shard_id}.manifest.parquet"
+
+                        if cache_local:
+                            import boto3
+
+                            s3_kwargs = {
+                                "aws_access_key_id": self.credentials[
+                                    "AWS_ACCESS_KEY_ID"
+                                ],
+                                "aws_secret_access_key": self.credentials[
+                                    "AWS_SECRET_ACCESS_KEY"
+                                ],
+                            }
+                            if "AWS_ENDPOINT" in self.credentials:
+                                s3_kwargs["endpoint_url"] = self.credentials[
+                                    "AWS_ENDPOINT"
+                                ]
+                            s3_client = boto3.client("s3", **s3_kwargs)
+
+                            bucket_name = Path(shard_id).parts[0]
+                            json_s3_key = str(Path(*Path(json_key).parts[1:]))
+                            parquet_s3_key = str(Path(*Path(parquet_key).parts[1:]))
+
+                            with open(json_path, "rb") as f:
+                                s3_client.put_object(
+                                    Bucket=bucket_name, Key=json_s3_key, Body=f.read()
+                                )
+                            with open(parquet_path, "rb") as f:
+                                s3_client.put_object(
+                                    Bucket=bucket_name,
+                                    Key=parquet_s3_key,
+                                    Body=f.read(),
+                                )
+                        else:
+                            self._fs.put(str(json_path), f"s3://{json_key}")
+                            self._fs.put(str(parquet_path), f"s3://{parquet_key}")
+
+            if len(written_shards) == 0:
+                raise RuntimeError("No shards written")
+
+            log.info("S3DataStore.write_shards: wrote %d shard(s)", len(written_shards))
 
         return written_shards
 

atdata/testing.py

@@ -14,7 +14,7 @@ Usage::
     samples = at_test.make_samples(MyType, n=100)
 
     # Use mock atmosphere client
-    client = at_test.MockAtmosphereClient()
+    client = at_test.MockAtmosphere()
 
     # Use in-memory index (SQLite backed, temporary)
     index = at_test.mock_index(tmp_path)
@@ -40,7 +40,7 @@ import webdataset as wds
 
 import atdata
 from atdata import Dataset, PackableSample
-from atdata.local._index import Index
+from atdata.index._index import Index
 from atdata.providers._sqlite import SqliteProvider
 
 ST = TypeVar("ST")
@@ -51,14 +51,14 @@ ST = TypeVar("ST")
 # ---------------------------------------------------------------------------
 
 
-class MockAtmosphereClient:
-    """In-memory mock of ``AtmosphereClient`` for testing.
+class MockAtmosphere:
+    """In-memory mock of ``Atmosphere`` for testing.
 
     Simulates login, schema publishing, dataset publishing, and record
     retrieval without requiring a live ATProto PDS.
 
     Examples:
-        >>> client = MockAtmosphereClient()
+        >>> client = MockAtmosphere()
         >>> client.login("alice.test", "password")
         >>> client.did
         'did:plc:mock000000000000'
@@ -294,8 +294,8 @@ try:
 
     @pytest.fixture
     def mock_atmosphere():
-        """Provide a fresh ``MockAtmosphereClient`` for each test."""
-        client = MockAtmosphereClient()
+        """Provide a fresh ``MockAtmosphere`` for each test."""
+        client = MockAtmosphere()
         client.login("test.mock.social", "test-password")
         yield client
         client.reset()
@@ -329,8 +329,12 @@ except ImportError:
 # Public API
 # ---------------------------------------------------------------------------
 
+# Deprecated alias for backward compatibility
+MockAtmosphereClient = MockAtmosphere
+
 __all__ = [
-    "MockAtmosphereClient",
+    "MockAtmosphere",
+    "MockAtmosphereClient",  # deprecated alias
     "make_dataset",
     "make_samples",
     "mock_index",

{atdata-0.3.0b1.dist-info → atdata-0.3.2b1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atdata
-Version: 0.3.0b1
+Version: 0.3.2b1
 Summary: A loose federation of distributed, typed datasets
 Author-email: Maxine Levesque <hello@maxine.science>, "Maxine @ Forecast Bio" <maxine@forecast.bio>
 License-File: LICENSE
@@ -30,7 +30,7 @@ Description-Content-Type: text/markdown
 
 # atdata
 
-[![codecov](https://codecov.io/gh/foundation-ac/atdata/branch/main/graph/badge.svg)](https://codecov.io/gh/foundation-ac/atdata)
+[![codecov](https://codecov.io/gh/forecast-bio/atdata/branch/main/graph/badge.svg)](https://codecov.io/gh/forecast-bio/atdata)
 
 A loose federation of distributed, typed datasets built on WebDataset.