atdata 0.2.3b1__py3-none-any.whl → 0.3.1b1__py3-none-any.whl

atdata/stores/_s3.py

@@ -0,0 +1,349 @@
+"""S3-compatible data store and helper functions."""
+
+from atdata import Dataset
+
+from pathlib import Path
+from uuid import uuid4
+from tempfile import TemporaryDirectory
+from dotenv import dotenv_values
+from typing import Any, BinaryIO, cast
+
+from s3fs import S3FileSystem
+import webdataset as wds
+
+
+def _s3_env(credentials_path: str | Path) -> dict[str, Any]:
+    """Load S3 credentials from .env file.
+
+    Args:
+        credentials_path: Path to .env file containing AWS_ENDPOINT,
+            AWS_ACCESS_KEY_ID, and AWS_SECRET_ACCESS_KEY.
+
+    Returns:
+        Dict with the three required credential keys.
+
+    Raises:
+        ValueError: If any required key is missing from the .env file.
+    """
+    credentials_path = Path(credentials_path)
+    env_values = dotenv_values(credentials_path)
+
+    required_keys = ("AWS_ENDPOINT", "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY")
+    missing = [k for k in required_keys if k not in env_values]
+    if missing:
+        raise ValueError(
+            f"Missing required keys in {credentials_path}: {', '.join(missing)}"
+        )
+
+    return {k: env_values[k] for k in required_keys}
+
+
+def _s3_from_credentials(creds: str | Path | dict) -> S3FileSystem:
+    """Create S3FileSystem from credentials dict or .env file path."""
+    if not isinstance(creds, dict):
+        creds = _s3_env(creds)
+
+    # Build kwargs, making endpoint_url optional
+    kwargs = {
+        "key": creds["AWS_ACCESS_KEY_ID"],
+        "secret": creds["AWS_SECRET_ACCESS_KEY"],
+    }
+    if "AWS_ENDPOINT" in creds:
+        kwargs["endpoint_url"] = creds["AWS_ENDPOINT"]
+
+    return S3FileSystem(**kwargs)
+
+
+def _create_s3_write_callbacks(
+    credentials: dict[str, Any],
+    temp_dir: str,
+    written_shards: list[str],
+    fs: S3FileSystem | None,
+    cache_local: bool,
+    add_s3_prefix: bool = False,
+) -> tuple:
+    """Create opener and post callbacks for ShardWriter with S3 upload.
+
+    Args:
+        credentials: S3 credentials dict.
+        temp_dir: Temporary directory for local caching.
+        written_shards: List to append written shard paths to.
+        fs: S3FileSystem for direct writes (used when cache_local=False).
+        cache_local: If True, write locally then copy to S3.
+        add_s3_prefix: If True, prepend 's3://' to shard paths.
+
+    Returns:
+        Tuple of (writer_opener, writer_post) callbacks.
+    """
+    if cache_local:
+        import boto3
+
+        s3_client_kwargs = {
+            "aws_access_key_id": credentials["AWS_ACCESS_KEY_ID"],
+            "aws_secret_access_key": credentials["AWS_SECRET_ACCESS_KEY"],
+        }
+        if "AWS_ENDPOINT" in credentials:
+            s3_client_kwargs["endpoint_url"] = credentials["AWS_ENDPOINT"]
+        s3_client = boto3.client("s3", **s3_client_kwargs)
+
+        def _writer_opener(p: str):
+            local_path = Path(temp_dir) / p
+            local_path.parent.mkdir(parents=True, exist_ok=True)
+            return open(local_path, "wb")
+
+        def _writer_post(p: str):
+            local_path = Path(temp_dir) / p
+            path_parts = Path(p).parts
+            bucket = path_parts[0]
+            key = str(Path(*path_parts[1:]))
+
+            with open(local_path, "rb") as f_in:
+                s3_client.put_object(Bucket=bucket, Key=key, Body=f_in.read())
+
+            local_path.unlink()
+            if add_s3_prefix:
+                written_shards.append(f"s3://{p}")
+            else:
+                written_shards.append(p)
+
+        return _writer_opener, _writer_post
+    else:
+        if fs is None:
+            raise ValueError("S3FileSystem required when cache_local=False")
+
+        def _direct_opener(s: str):
+            return cast(BinaryIO, fs.open(f"s3://{s}", "wb"))
+
+        def _direct_post(s: str):
+            if add_s3_prefix:
+                written_shards.append(f"s3://{s}")
+            else:
+                written_shards.append(s)
+
+        return _direct_opener, _direct_post
+
+
+class S3DataStore:
+    """S3-compatible data store implementing AbstractDataStore protocol.
+
+    Handles writing dataset shards to S3-compatible object storage and
+    resolving URLs for reading.
+
+    Attributes:
+        credentials: S3 credentials dictionary.
+        bucket: Target bucket name.
+        _fs: S3FileSystem instance.
+    """
+
+    def __init__(
+        self,
+        credentials: str | Path | dict[str, Any],
+        *,
+        bucket: str,
+    ) -> None:
+        """Initialize an S3 data store.
+
+        Args:
+            credentials: Path to .env file or dict with AWS_ACCESS_KEY_ID,
+                AWS_SECRET_ACCESS_KEY, and optionally AWS_ENDPOINT.
+            bucket: Name of the S3 bucket for storage.
+        """
+        if isinstance(credentials, dict):
+            self.credentials = credentials
+        else:
+            self.credentials = _s3_env(credentials)
+
+        self.bucket = bucket
+        self._fs = _s3_from_credentials(self.credentials)
+
+    def write_shards(
+        self,
+        ds: Dataset,
+        *,
+        prefix: str,
+        cache_local: bool = False,
+        manifest: bool = False,
+        schema_version: str = "1.0.0",
+        source_job_id: str | None = None,
+        parent_shards: list[str] | None = None,
+        pipeline_version: str | None = None,
+        **kwargs,
+    ) -> list[str]:
+        """Write dataset shards to S3.
+
+        Args:
+            ds: The Dataset to write.
+            prefix: Path prefix within bucket (e.g., 'datasets/mnist/v1').
+            cache_local: If True, write locally first then copy to S3.
+            manifest: If True, generate per-shard manifest files alongside
+                each tar shard (``.manifest.json`` + ``.manifest.parquet``).
+            schema_version: Schema version for manifest headers.
+            source_job_id: Optional provenance job identifier for manifests.
+            parent_shards: Optional list of input shard identifiers for provenance.
+            pipeline_version: Optional pipeline version string for provenance.
+            **kwargs: Additional args passed to wds.ShardWriter (e.g., maxcount).
+
+        Returns:
+            List of S3 URLs for the written shards.
+
+        Raises:
+            RuntimeError: If no shards were written.
+        """
+        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,
+            )
+
+            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,
+                        )
+                        # 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()
+                            )
+                        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")
+
+        return written_shards
+
+    def read_url(self, url: str) -> str:
+        """Resolve an S3 URL for reading/streaming.
+
+        For S3-compatible stores with custom endpoints (like Cloudflare R2,
+        MinIO, etc.), converts s3:// URLs to HTTPS URLs that WebDataset can
+        stream directly.
+
+        For standard AWS S3 (no custom endpoint), URLs are returned unchanged
+        since WebDataset's built-in s3fs integration handles them.
+
+        Args:
+            url: S3 URL to resolve (e.g., 's3://bucket/path/file.tar').
+
+        Returns:
+            HTTPS URL if custom endpoint is configured, otherwise unchanged.
+            Example: 's3://bucket/path' -> 'https://endpoint.com/bucket/path'
+        """
+        endpoint = self.credentials.get("AWS_ENDPOINT")
+        if endpoint and url.startswith("s3://"):
+            # s3://bucket/path -> https://endpoint/bucket/path
+            path = url[5:]  # Remove 's3://' prefix
+            endpoint = endpoint.rstrip("/")
+            return f"{endpoint}/{path}"
+        return url
+
+    def supports_streaming(self) -> bool:
+        """S3 supports streaming reads.
+
+        Returns:
+            True.
+        """
+        return True

atdata/testing.py

@@ -0,0 +1,341 @@
+"""Testing utilities for atdata.
+
+Provides mock clients, dataset factories, and pytest fixtures for writing
+tests against atdata without requiring external services (Redis, S3, ATProto PDS).
+
+Usage::
+
+    import atdata.testing as at_test
+
+    # Create a dataset from samples
+    ds = at_test.make_dataset(tmp_path, [sample1, sample2])
+
+    # Generate random samples
+    samples = at_test.make_samples(MyType, n=100)
+
+    # Use mock atmosphere client
+    client = at_test.MockAtmosphere()
+
+    # Use in-memory index (SQLite backed, temporary)
+    index = at_test.mock_index(tmp_path)
+
+Pytest fixtures (available when ``atdata`` is installed)::
+
+    def test_something(mock_atmosphere):
+        client = mock_atmosphere
+        client.login("user", "pass")
+        ...
+"""
+
+from __future__ import annotations
+
+import tempfile
+import uuid
+from dataclasses import fields as dc_fields
+from pathlib import Path
+from typing import Any, Sequence, Type, TypeVar
+
+import numpy as np
+import webdataset as wds
+
+import atdata
+from atdata import Dataset, PackableSample
+from atdata.index._index import Index
+from atdata.providers._sqlite import SqliteProvider
+
+ST = TypeVar("ST")
+
+
+# ---------------------------------------------------------------------------
+# Mock Atmosphere Client
+# ---------------------------------------------------------------------------
+
+
+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 = MockAtmosphere()
+        >>> client.login("alice.test", "password")
+        >>> client.did
+        'did:plc:mock000000000000'
+    """
+
+    def __init__(
+        self,
+        did: str = "did:plc:mock000000000000",
+        handle: str = "test.mock.social",
+    ) -> None:
+        self.did = did
+        self.handle = handle
+        self._logged_in = False
+        self._records: dict[str, dict[str, Any]] = {}
+        self._schemas: dict[str, dict[str, Any]] = {}
+        self._datasets: dict[str, dict[str, Any]] = {}
+        self._blobs: dict[str, bytes] = {}
+        self._session_string = "mock-session-string"
+        self._call_log: list[tuple[str, dict[str, Any]]] = []
+
+    def login(self, handle: str, password: str) -> dict[str, Any]:
+        """Simulate login. Always succeeds."""
+        self._logged_in = True
+        self.handle = handle
+        self._call_log.append(("login", {"handle": handle}))
+        return {"did": self.did, "handle": self.handle}
+
+    @property
+    def is_authenticated(self) -> bool:
+        return self._logged_in
+
+    def export_session_string(self) -> str:
+        return self._session_string
+
+    def create_record(
+        self,
+        collection: str,
+        record: dict[str, Any],
+        rkey: str | None = None,
+    ) -> str:
+        """Simulate creating a record. Returns a mock AT URI."""
+        key = rkey or uuid.uuid4().hex[:12]
+        uri = f"at://{self.did}/{collection}/{key}"
+        self._records[uri] = record
+        self._call_log.append(
+            ("create_record", {"collection": collection, "rkey": key, "uri": uri})
+        )
+        return uri
+
+    def get_record(self, uri: str) -> dict[str, Any]:
+        """Retrieve a previously created record by URI."""
+        if uri not in self._records:
+            raise KeyError(f"Record not found: {uri}")
+        return self._records[uri]
+
+    def list_records(self, collection: str) -> list[dict[str, Any]]:
+        """List records for a collection."""
+        return [
+            {"uri": uri, "value": rec}
+            for uri, rec in self._records.items()
+            if collection in uri
+        ]
+
+    def upload_blob(self, data: bytes) -> dict[str, Any]:
+        """Simulate uploading a blob. Returns a mock blob ref."""
+        ref = f"blob:{uuid.uuid4().hex[:16]}"
+        self._blobs[ref] = data
+        self._call_log.append(("upload_blob", {"ref": ref, "size": len(data)}))
+        return {"ref": {"$link": ref}, "mimeType": "application/octet-stream"}
+
+    def get_blob(self, did: str, cid: str) -> bytes:
+        """Retrieve a previously uploaded blob."""
+        if cid not in self._blobs:
+            raise KeyError(f"Blob not found: {cid}")
+        return self._blobs[cid]
+
+    def reset(self) -> None:
+        """Clear all stored state."""
+        self._records.clear()
+        self._schemas.clear()
+        self._datasets.clear()
+        self._blobs.clear()
+        self._call_log.clear()
+        self._logged_in = False
+
+
+# ---------------------------------------------------------------------------
+# Dataset Factory
+# ---------------------------------------------------------------------------
+
+
+def make_dataset(
+    path: Path,
+    samples: Sequence[PackableSample],
+    *,
+    name: str = "test",
+    sample_type: type | None = None,
+) -> Dataset:
+    """Create a ``Dataset`` from a list of samples.
+
+    Writes the samples to a WebDataset tar file in *path* and returns a
+    ``Dataset`` configured to read them back.
+
+    Args:
+        path: Directory where the tar file will be created.
+        samples: Sequence of ``PackableSample`` (or ``@packable``) instances.
+        name: Filename prefix for the tar file.
+        sample_type: Explicit sample type for the Dataset generic parameter.
+            If ``None``, inferred from the first sample.
+
+    Returns:
+        A ``Dataset`` ready for iteration.
+
+    Examples:
+        >>> ds = make_dataset(tmp_path, [MySample(x=1), MySample(x=2)])
+        >>> assert len(list(ds.ordered())) == 2
+    """
+    if not samples:
+        raise ValueError("samples must be non-empty")
+
+    tar_path = path / f"{name}-000000.tar"
+    tar_path.parent.mkdir(parents=True, exist_ok=True)
+
+    with wds.writer.TarWriter(str(tar_path)) as writer:
+        for sample in samples:
+            writer.write(sample.as_wds)
+
+    st = sample_type or type(samples[0])
+    return Dataset[st](url=str(tar_path))
+
+
+def make_samples(
+    sample_type: Type[ST], n: int = 10, seed: int | None = None
+) -> list[ST]:
+    """Generate *n* random instances of a ``@packable`` sample type.
+
+    Inspects the dataclass fields and generates appropriate random data:
+    - ``str`` fields get ``"field_name_0"``, ``"field_name_1"``, etc.
+    - ``int`` fields get sequential integers
+    - ``float`` fields get random floats in [0, 1)
+    - ``bool`` fields alternate True/False
+    - ``bytes`` fields get random 16 bytes
+    - NDArray fields get random ``(4, 4)`` float32 arrays
+
+    Args:
+        sample_type: A ``@packable``-decorated class or ``PackableSample`` subclass.
+        n: Number of samples to generate.
+        seed: Optional random seed for reproducibility.
+
+    Returns:
+        List of *n* sample instances.
+
+    Examples:
+        >>> @atdata.packable
+        ... class Point:
+        ...     x: float
+        ...     y: float
+        ...     label: str
+        >>> points = make_samples(Point, n=5, seed=42)
+        >>> len(points)
+        5
+    """
+    rng = np.random.default_rng(seed)
+    result: list[ST] = []
+
+    for i in range(n):
+        kwargs: dict[str, Any] = {}
+        for field in dc_fields(sample_type):
+            type_str = str(field.type)
+            fname = field.name
+
+            if field.type is str or type_str == "str":
+                kwargs[fname] = f"{fname}_{i}"
+            elif field.type is int or type_str == "int":
+                kwargs[fname] = i
+            elif field.type is float or type_str == "float":
+                kwargs[fname] = float(rng.random())
+            elif field.type is bool or type_str == "bool":
+                kwargs[fname] = i % 2 == 0
+            elif field.type is bytes or type_str == "bytes":
+                kwargs[fname] = rng.bytes(16)
+            elif "NDArray" in type_str or "ndarray" in type_str.lower():
+                kwargs[fname] = rng.standard_normal((4, 4)).astype(np.float32)
+            elif "list" in type_str.lower():
+                kwargs[fname] = [f"{fname}_{i}_{j}" for j in range(3)]
+            elif "None" in type_str:
+                # Optional field — leave at default
+                if field.default is not field.default_factory:  # type: ignore[attr-defined]
+                    continue
+            else:
+                kwargs[fname] = f"{fname}_{i}"
+
+        result.append(sample_type(**kwargs))
+
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Mock Index
+# ---------------------------------------------------------------------------
+
+
+def mock_index(path: Path | None = None, **kwargs: Any) -> Index:
+    """Create an in-memory SQLite-backed ``Index`` for testing.
+
+    No Redis or external services required.
+
+    Args:
+        path: Directory for the SQLite database file. If ``None``, uses
+            a temporary directory.
+        **kwargs: Additional keyword arguments passed to ``Index()``.
+
+    Returns:
+        An ``Index`` instance backed by a temporary SQLite database.
+
+    Examples:
+        >>> index = mock_index(tmp_path)
+        >>> ref = index.publish_schema(MyType, version="1.0.0")
+    """
+    if path is None:
+        path = Path(tempfile.mkdtemp())
+    db_path = path / "test_index.db"
+    provider = SqliteProvider(str(db_path))
+    return Index(provider=provider, atmosphere=None, **kwargs)
+
+
+# ---------------------------------------------------------------------------
+# Pytest plugin (fixtures auto-discovered when atdata is installed)
+# ---------------------------------------------------------------------------
+
+try:
+    import pytest
+
+    @pytest.fixture
+    def mock_atmosphere():
+        """Provide a fresh ``MockAtmosphere`` for each test."""
+        client = MockAtmosphere()
+        client.login("test.mock.social", "test-password")
+        yield client
+        client.reset()
+
+    @pytest.fixture
+    def tmp_dataset(tmp_path: Path):
+        """Provide a small ``Dataset[SharedBasicSample]`` with 10 samples.
+
+        Uses ``SharedBasicSample`` (name: str, value: int) from the test suite.
+        """
+
+        @atdata.packable
+        class _TmpSample:
+            name: str
+            value: int
+
+        samples = [_TmpSample(name=f"s{i}", value=i) for i in range(10)]
+        return make_dataset(tmp_path, samples, sample_type=_TmpSample)
+
+    @pytest.fixture
+    def tmp_index(tmp_path: Path):
+        """Provide a fresh SQLite-backed ``Index`` for each test."""
+        return mock_index(tmp_path)
+
+except ImportError:
+    # pytest not installed — skip fixture registration
+    _no_pytest = True
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+# Deprecated alias for backward compatibility
+MockAtmosphereClient = MockAtmosphere
+
+__all__ = [
+    "MockAtmosphere",
+    "MockAtmosphereClient",  # deprecated alias
+    "make_dataset",
+    "make_samples",
+    "mock_index",
+]