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

atdata/testing.py

@@ -0,0 +1,337 @@
+"""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.MockAtmosphereClient()
+
+    # 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.local._index import Index
+from atdata.providers._sqlite import SqliteProvider
+
+ST = TypeVar("ST")
+
+
+# ---------------------------------------------------------------------------
+# Mock Atmosphere Client
+# ---------------------------------------------------------------------------
+
+
+class MockAtmosphereClient:
+    """In-memory mock of ``AtmosphereClient`` for testing.
+
+    Simulates login, schema publishing, dataset publishing, and record
+    retrieval without requiring a live ATProto PDS.
+
+    Examples:
+        >>> client = MockAtmosphereClient()
+        >>> 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 ``MockAtmosphereClient`` for each test."""
+        client = MockAtmosphereClient()
+        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
+# ---------------------------------------------------------------------------
+
+__all__ = [
+    "MockAtmosphereClient",
+    "make_dataset",
+    "make_samples",
+    "mock_index",
+]

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atdata
-Version: 0.2.3b1
+Version: 0.3.0b1
 Summary: A loose federation of distributed, typed datasets
 Author-email: Maxine Levesque <hello@maxine.science>, "Maxine @ Forecast Bio" <maxine@forecast.bio>
 License-File: LICENSE
@@ -20,9 +20,12 @@ Requires-Dist: requests>=2.32.5
 Requires-Dist: s3fs>=2025.12.0
 Requires-Dist: schemamodels>=0.9.1
 Requires-Dist: tqdm>=4.67.1
+Requires-Dist: typer>=0.21.1
 Requires-Dist: webdataset>=1.0.2
 Provides-Extra: atmosphere
 Requires-Dist: atproto>=0.0.55; extra == 'atmosphere'
+Provides-Extra: postgres
+Requires-Dist: psycopg[binary]>=3.1; extra == 'postgres'
 Description-Content-Type: text/markdown
 
 # atdata

atdata-0.3.0b1.dist-info/RECORD

@@ -0,0 +1,54 @@
+atdata/.gitignore,sha256=ROcLwaiURIGOAYGFpakac_WEeKS4hH4HxJuscfDGo2s,10
+atdata/__init__.py,sha256=TXEvc6R_VQ4zzjHr9OIqi_hhwfc6OqerVWajtDRwJTA,3249
+atdata/_cid.py,sha256=6wLV_dcQJy5Eb-wld7_h7Kcp7QoVixIqUDIIoSwpQms,3992
+atdata/_exceptions.py,sha256=a2IzmTJqLyIgQ0RkW3KhaaaSDdQJGLoXzuHfEX0DfB4,5390
+atdata/_helpers.py,sha256=e-1lChOY9jvihs5Q04Tp5Rp9oKIoVf5N4wuuwv0Hl64,2224
+atdata/_hf_api.py,sha256=ThNoVDIztbJCLtIKSCol4Tm7Lf1TqbU4E-fFsKt5FaA,24318
+atdata/_logging.py,sha256=Wj9MTe9xCP_awiQYfa7219Fb443Rkwu_jIqM-m2u2GA,2225
+atdata/_protocols.py,sha256=LYeM0eIn5nQ56DssdGFd-CvJFXh9C3YR2GLOSFxg1Kw,14479
+atdata/_schema_codec.py,sha256=OpRs8bfCPGCEJsBKlp8rfndBldco3jwQQ7k9npN-EPQ,14395
+atdata/_sources.py,sha256=A7HMkS_dqN5Sx7rG1nZsO1Laxozt3C65_P2Hiv41VXk,16624
+atdata/_stub_manager.py,sha256=Heh0HAYjVjnkUQcPWAEOrkkEKN3Mi9vTJPA-ZRseFw8,19141
+atdata/_type_utils.py,sha256=mhD0ywXgIEpVaAWpGVKGneUOh0cmI96zKz6g2O1bdg0,3762
+atdata/dataset.py,sha256=ogclAFyeDh8ZGvgiDViiDn_sbkh94rpXgR-jtwFpi6M,41123
+atdata/lens.py,sha256=s6Ebeacth4cF_yCmnH4MuZ6TwQJmUmrFZ3eIh1jQSZI,9914
+atdata/promote.py,sha256=W8I_OCTStKY_Mu1OJvtVyJrWqmXjRQaTQLJ_Id-nBmA,6274
+atdata/repository.py,sha256=5ZlYOw5VQdkuN0rntfw26CA_O_YsZLYMSGaRPPPn9Gc,10205
+atdata/testing.py,sha256=OUPXn8CBuxNVyhj-lgIiw_YkTNrewdEkgFyfwogLH3E,10858
+atdata/atmosphere/__init__.py,sha256=g8larX6DIA4AVFgt-g3vm6Ijsb4vU3YOAFPNCDMd-sk,10064
+atdata/atmosphere/_types.py,sha256=MRhXnmAuQLJPdoq1skrBGXCsaQYdtKG_nA3YlSjwJXY,9595
+atdata/atmosphere/client.py,sha256=acw82w3_cxWbWDtIRvH1VDHGJSroGqhSenFFostXTXo,16210
+atdata/atmosphere/lens.py,sha256=EnrddTD-SAnyxU0XF__QkePLUhb4lPwfFLaseL_STDc,9260
+atdata/atmosphere/records.py,sha256=Vkc7K0XFozwXVaDgXP706Z1-De1C95ZPtra1EQE3MGQ,15860
+atdata/atmosphere/schema.py,sha256=xvIyJbzbZrcAi7TOWFQXXfgqAyuhb5Yqttf9dSy9x50,7758
+atdata/atmosphere/store.py,sha256=NR4tGS9u3_ogvnyyOHDVF0tRKChruj_NE9Df4qrZiDU,6324
+atdata/cli/__init__.py,sha256=3ZtMcYmx78L_8LqzPHaLlEOlFXl_uHBJtffSL6UB-kE,5834
+atdata/cli/diagnose.py,sha256=Det9ozOvxXKd8Abu-xEsMjaXR34H_cuSX9MJJIlhnsA,5483
+atdata/cli/inspect.py,sha256=YoJkWeXj1cHcOuDH9Z3btWMhiFNOG6JG3x-tScee-84,2019
+atdata/cli/local.py,sha256=dYosOuVmp_PIM8re4DFsm00uspUJftTOZ0TBBNlTqAM,8105
+atdata/cli/preview.py,sha256=n90Eqj920JEi7nBl4NnbmEhZxKLzHoR2CfNI9lkbGuc,1735
+atdata/cli/schema.py,sha256=JwbQOBOPcvDTF6QgptuqZF5VEEMHrjxlj959q6E1DQY,2882
+atdata/local/__init__.py,sha256=ZBJspq-dKOYZESUfxnQ0vIqpyBXGo2Fy116RvfFtDgs,1872
+atdata/local/_entry.py,sha256=Cea1uAbURlBVn3vlrcLD5mtwIaqYoHjgb1GVsBqrszw,5264
+atdata/local/_index.py,sha256=lCl0JAmHxmhz9lFlMiMvO_TApYARK3dyiBAJTVUaYO4,34305
+atdata/local/_repo_legacy.py,sha256=7njQTp9HReyWNcIOyYG_Zg90nhEI8KRVCLI0ANnoUpY,7541
+atdata/local/_s3.py,sha256=7KuIIZgqaiaNjULlKC3QPbLkKaj7rRVSYJt5zU_OEXI,12809
+atdata/local/_schema.py,sha256=myCZ7eTJ0hKLQuF5qBvjzttyenKSqXd3YDdamFdRw1E,11634
+atdata/manifest/__init__.py,sha256=GB7V8OXLDAIuUX4JSBVL7_UV9CFALMiMfoeZvoNkug0,1309
+atdata/manifest/_aggregates.py,sha256=yabfBJtdVqOBWbkNqg9gp0BgpG62wi3QsGYlHgAAcFk,4198
+atdata/manifest/_builder.py,sha256=abU5cnRnIJZ_qf_BKr6-fMYU1HaJWkBVgmGL0WzRRzU,5252
+atdata/manifest/_fields.py,sha256=4Mm31HGb3qRwvp5Kn9Pp-4dX0tlNrQHxEDZJnBzFpos,5003
+atdata/manifest/_manifest.py,sha256=hPcURPjjPS5r3H_veVshwT6Qw-lC9eQll84Fos1xrCE,4852
+atdata/manifest/_query.py,sha256=25_0i60jWQhuMXsuuVLvSPS8v9pejTi5Bik99e2L99A,4641
+atdata/manifest/_writer.py,sha256=_g2YZof-GnN3nVLpAAE5fDmsV9GzlPYf2oa1BWYswxI,2206
+atdata/providers/__init__.py,sha256=LDziQc6Dc0QCfJcr_sDDqxArstDB0dSZHGgXCIIXbzo,848
+atdata/providers/_base.py,sha256=3ICfgo2AWrRyAHrH8AJM4F0ssDBVD_K5iZAuQcUFLGA,4402
+atdata/providers/_factory.py,sha256=9clKbch3gxqJWufY-HIbbfaBstzoq1LOcMrJOKHTPhw,2039
+atdata/providers/_postgres.py,sha256=BpbjiD9jifLO74aWttSq4JgVoecAypCoYzO_7DdAZuQ,7402
+atdata/providers/_redis.py,sha256=VljNKSsLAX5H0NoOuqkuLQeQAFuuettq1MnpBJgnqhw,6054
+atdata/providers/_sqlite.py,sha256=urZWl7ErgKJJ1uBpnmmdkz7uWnrKGkJDIaTvAdyN7Uo,6486
+atdata-0.3.0b1.dist-info/METADATA,sha256=zxCV6U6eMrtYZII4hBIrY7tZ9hs4ntTdAwPRhg_HISk,7410
+atdata-0.3.0b1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+atdata-0.3.0b1.dist-info/entry_points.txt,sha256=6-iQr1veSTq-ac94bLyfcyGHprrZWevPEd12BWX37tQ,39
+atdata-0.3.0b1.dist-info/licenses/LICENSE,sha256=Pz2eACSxkhsGfW9_iN60pgy-enjnbGTj8df8O3ebnQQ,16726
+atdata-0.3.0b1.dist-info/RECORD,,