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

atdata/lexicons/ac.foundation.dataset.lens.json

@@ -0,0 +1,99 @@
+{
+  "lexicon": 1,
+  "id": "ac.foundation.dataset.lens",
+  "defs": {
+    "main": {
+      "type": "record",
+      "description": "Bidirectional transformation (Lens) between two sample types, with code stored in external repositories",
+      "key": "tid",
+      "record": {
+        "type": "object",
+        "required": [
+          "name",
+          "sourceSchema",
+          "targetSchema",
+          "getterCode",
+          "putterCode",
+          "createdAt"
+        ],
+        "properties": {
+          "name": {
+            "type": "string",
+            "description": "Human-readable lens name",
+            "maxLength": 100
+          },
+          "sourceSchema": {
+            "type": "string",
+            "description": "AT-URI reference to source schema",
+            "maxLength": 500
+          },
+          "targetSchema": {
+            "type": "string",
+            "description": "AT-URI reference to target schema",
+            "maxLength": 500
+          },
+          "description": {
+            "type": "string",
+            "description": "What this transformation does",
+            "maxLength": 1000
+          },
+          "getterCode": {
+            "type": "ref",
+            "ref": "#codeReference",
+            "description": "Code reference for getter function (Source -> Target)"
+          },
+          "putterCode": {
+            "type": "ref",
+            "ref": "#codeReference",
+            "description": "Code reference for putter function (Target, Source -> Source)"
+          },
+          "language": {
+            "type": "string",
+            "description": "Programming language of the lens implementation (e.g., 'python', 'typescript')",
+            "maxLength": 50
+          },
+          "metadata": {
+            "type": "object",
+            "description": "Arbitrary metadata (author, performance notes, etc.)"
+          },
+          "createdAt": {
+            "type": "string",
+            "format": "datetime",
+            "description": "Timestamp when this lens was created"
+          }
+        }
+      }
+    },
+    "codeReference": {
+      "type": "object",
+      "description": "Reference to code in an external repository (GitHub, tangled.org, etc.)",
+      "required": [
+        "repository",
+        "commit",
+        "path"
+      ],
+      "properties": {
+        "repository": {
+          "type": "string",
+          "description": "Repository URL (e.g., 'https://github.com/user/repo' or 'at://did/tangled.repo/...')",
+          "maxLength": 500
+        },
+        "commit": {
+          "type": "string",
+          "description": "Git commit hash (ensures immutability)",
+          "maxLength": 40
+        },
+        "path": {
+          "type": "string",
+          "description": "Path to function within repository (e.g., 'lenses/vision.py:rgb_to_grayscale')",
+          "maxLength": 500
+        },
+        "branch": {
+          "type": "string",
+          "description": "Optional branch name (for reference, commit hash is authoritative)",
+          "maxLength": 100
+        }
+      }
+    }
+  }
+}

atdata/lexicons/ac.foundation.dataset.record.json

@@ -0,0 +1,96 @@
+{
+  "lexicon": 1,
+  "id": "ac.foundation.dataset.record",
+  "defs": {
+    "main": {
+      "type": "record",
+      "description": "Index record for a WebDataset-backed dataset with references to storage location and sample schema",
+      "key": "tid",
+      "record": {
+        "type": "object",
+        "required": [
+          "name",
+          "schemaRef",
+          "storage",
+          "createdAt"
+        ],
+        "properties": {
+          "name": {
+            "type": "string",
+            "description": "Human-readable dataset name",
+            "maxLength": 200
+          },
+          "schemaRef": {
+            "type": "string",
+            "format": "at-uri",
+            "description": "AT-URI reference to the schema record for this dataset's samples",
+            "maxLength": 500
+          },
+          "storage": {
+            "type": "union",
+            "description": "Storage location for dataset files (WebDataset tar archives)",
+            "refs": [
+              "ac.foundation.dataset.storageExternal",
+              "ac.foundation.dataset.storageBlobs"
+            ]
+          },
+          "description": {
+            "type": "string",
+            "description": "Human-readable description of the dataset",
+            "maxLength": 5000
+          },
+          "metadata": {
+            "type": "bytes",
+            "description": "Msgpack-encoded metadata dict for arbitrary extended key-value pairs. Use this for additional metadata beyond the core top-level fields (license, tags, size). Top-level fields are preferred for discoverable/searchable metadata.",
+            "maxLength": 100000
+          },
+          "tags": {
+            "type": "array",
+            "description": "Searchable tags for dataset discovery. Aligns with Schema.org keywords property.",
+            "items": {
+              "type": "string",
+              "maxLength": 150
+            },
+            "maxLength": 30
+          },
+          "size": {
+            "type": "ref",
+            "ref": "#datasetSize",
+            "description": "Dataset size information (optional)"
+          },
+          "license": {
+            "type": "string",
+            "description": "License identifier or URL. SPDX identifiers recommended (e.g., MIT, Apache-2.0, CC-BY-4.0) or full SPDX URLs (e.g., http://spdx.org/licenses/MIT). Aligns with Schema.org license property.",
+            "maxLength": 200
+          },
+          "createdAt": {
+            "type": "string",
+            "format": "datetime",
+            "description": "Timestamp when this dataset record was created"
+          }
+        }
+      }
+    },
+    "datasetSize": {
+      "type": "object",
+      "description": "Information about dataset size",
+      "properties": {
+        "samples": {
+          "type": "integer",
+          "description": "Total number of samples in the dataset",
+          "minimum": 0
+        },
+        "bytes": {
+          "type": "integer",
+          "description": "Total size in bytes",
+          "minimum": 0
+        },
+        "shards": {
+          "type": "integer",
+          "description": "Number of WebDataset shards",
+          "minimum": 1
+        }
+      }
+    }
+  }
+}

atdata/lexicons/ac.foundation.dataset.schema.json

@@ -0,0 +1,107 @@
+{
+  "lexicon": 1,
+  "id": "ac.foundation.dataset.schema",
+  "defs": {
+    "main": {
+      "type": "record",
+      "description": "Definition of a PackableSample-compatible sample type. Supports versioning via rkey format: {NSID}@{semver}. Schema format is extensible via union type.",
+      "key": "any",
+      "record": {
+        "type": "object",
+        "required": [
+          "name",
+          "version",
+          "schemaType",
+          "schema",
+          "createdAt"
+        ],
+        "properties": {
+          "name": {
+            "type": "string",
+            "description": "Human-readable display name for this sample type. Used for documentation and UI. The NSID in the record URI provides unique identification; name collisions across NSIDs are acceptable.",
+            "maxLength": 100
+          },
+          "version": {
+            "type": "string",
+            "description": "Semantic version (e.g., '1.0.0')",
+            "pattern": "^(0|[1-9]\\d*)\\.(0|[1-9]\\d*)\\.(0|[1-9]\\d*)(?:-((?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\\.(?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\\+([0-9a-zA-Z-]+(?:\\.[0-9a-zA-Z-]+)*))?$",
+            "maxLength": 100
+          },
+          "schemaType": {
+            "type": "ref",
+            "ref": "ac.foundation.dataset.schemaType",
+            "description": "Type of schema definition. This field indicates which union member is present in the schema field."
+          },
+          "schema": {
+            "type": "union",
+            "refs": ["ac.foundation.dataset.schema#jsonSchemaFormat"],
+            "closed": false,
+            "description": "Schema definition for this sample type. Currently supports JSON Schema Draft 7. Union allows for future schema formats (Avro, Protobuf, etc.) without breaking changes."
+          },
+          "description": {
+            "type": "string",
+            "description": "Human-readable description of what this sample type represents",
+            "maxLength": 5000
+          },
+          "metadata": {
+            "type": "object",
+            "description": "Optional metadata about this schema. Common fields include license and tags, but any additional fields are permitted.",
+            "maxProperties": 50,
+            "properties": {
+              "license": {
+                "type": "string",
+                "description": "License identifier or URL. SPDX identifiers recommended (e.g., MIT, Apache-2.0, CC-BY-4.0) or full SPDX URLs (e.g., http://spdx.org/licenses/MIT). Aligns with Schema.org license property.",
+                "maxLength": 200
+              },
+              "tags": {
+                "type": "array",
+                "description": "Categorization keywords for discovery. Aligns with Schema.org keywords property.",
+                "items": {
+                  "type": "string",
+                  "maxLength": 150
+                },
+                "maxLength": 30
+              }
+            }
+          },
+          "createdAt": {
+            "type": "string",
+            "format": "datetime",
+            "description": "Timestamp when this schema version was created. Immutable once set (ATProto records are permanent)."
+          }
+        }
+      }
+    },
+    "jsonSchemaFormat": {
+      "type": "object",
+      "description": "JSON Schema Draft 7 format for sample type definitions. Used with NDArray shim for array types.",
+      "required": ["$type", "$schema", "type", "properties"],
+      "properties": {
+        "$type": {
+          "type": "string",
+          "const": "ac.foundation.dataset.schema#jsonSchemaFormat"
+        },
+        "$schema": {
+          "type": "string",
+          "const": "http://json-schema.org/draft-07/schema#",
+          "description": "JSON Schema version identifier"
+        },
+        "type": {
+          "type": "string",
+          "const": "object",
+          "description": "Sample types must be objects"
+        },
+        "properties": {
+          "type": "object",
+          "description": "Field definitions for the sample type",
+          "minProperties": 1
+        },
+        "arrayFormatVersions": {
+          "type": "object",
+          "description": "Mapping from array format identifiers to semantic versions. Keys are ac.foundation.dataset.arrayFormat values (e.g., 'ndarrayBytes'), values are semver strings (e.g., '1.0.0'). Foundation.ac maintains canonical shim schemas at https://foundation.ac/schemas/atdata-{format}-bytes/{version}/.",
+          "maxProperties": 10
+        }
+      }
+    }
+  }
+}

atdata/lexicons/ac.foundation.dataset.schemaType.json

@@ -0,0 +1,16 @@
+{
+  "lexicon": 1,
+  "id": "ac.foundation.dataset.schemaType",
+  "defs": {
+    "main": {
+      "type": "string",
+      "description": "Schema type identifier for atdata sample definitions. Known values correspond to token definitions in this Lexicon. New schema types can be added as tokens without breaking changes.",
+      "knownValues": ["jsonSchema"],
+      "maxLength": 50
+    },
+    "jsonSchema": {
+      "type": "token",
+      "description": "JSON Schema Draft 7 format for sample type definitions. When schemaType is 'jsonSchema', the schema field must contain an object conforming to ac.foundation.dataset.schema#jsonSchemaFormat."
+    }
+  }
+}

atdata/lexicons/ac.foundation.dataset.storageBlobs.json

@@ -0,0 +1,24 @@
+{
+  "lexicon": 1,
+  "id": "ac.foundation.dataset.storageBlobs",
+  "defs": {
+    "main": {
+      "type": "object",
+      "description": "Storage via ATProto PDS blobs for WebDataset tar archives. Each blob contains one or more tar files. Used in ac.foundation.dataset.record storage union for maximum decentralization.",
+      "required": [
+        "blobs"
+      ],
+      "properties": {
+        "blobs": {
+          "type": "array",
+          "description": "Array of blob references for WebDataset tar files",
+          "items": {
+            "type": "blob",
+            "description": "Blob reference to a WebDataset tar archive"
+          },
+          "minLength": 1
+        }
+      }
+    }
+  }
+}

atdata/lexicons/ac.foundation.dataset.storageExternal.json

@@ -0,0 +1,25 @@
+{
+  "lexicon": 1,
+  "id": "ac.foundation.dataset.storageExternal",
+  "defs": {
+    "main": {
+      "type": "object",
+      "description": "External storage via URLs (S3, HTTP, IPFS, etc.) for WebDataset tar archives. URLs support brace notation for sharding (e.g., 'data-{000000..000099}.tar'). Used in ac.foundation.dataset.record storage union.",
+      "required": [
+        "urls"
+      ],
+      "properties": {
+        "urls": {
+          "type": "array",
+          "description": "WebDataset URLs with optional brace notation for sharded tar files",
+          "items": {
+            "type": "string",
+            "format": "uri",
+            "maxLength": 1000
+          },
+          "minLength": 1
+        }
+      }
+    }
+  }
+}

atdata/lexicons/ndarray_shim.json

@@ -0,0 +1,16 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "$id": "https://foundation.ac/schemas/atdata-ndarray-bytes/1.0.0",
+  "title": "ATDataNDArrayBytes",
+  "description": "Standard definition for numpy NDArray types in JSON Schema, compatible with atdata WebDataset serialization. This type's contents are interpreted as containing the raw bytes data for a serialized numpy NDArray, and serve as a marker for atdata-based code generation to use standard numpy types, rather than generated dataclasses.",
+  "version": "1.0.0",
+  "$defs": {
+    "ndarray": {
+      "type": "string",
+      "format": "byte",
+      "description": "Numpy array serialized using numpy `.npy` format via `np.save` (includes dtype and shape in binary header). When represented in JSON, this is a base64-encoded string. In msgpack, this is raw bytes.",
+      "contentEncoding": "base64",
+      "contentMediaType": "application/octet-stream"
+    }
+  }
+}

atdata/local/__init__.py

@@ -0,0 +1,70 @@
+"""Backward-compatibility shim for atdata.local.
+
+.. deprecated::
+    Import from ``atdata.index`` and ``atdata.stores`` instead::
+
+        from atdata.index import Index, LocalDatasetEntry
+        from atdata.stores import S3DataStore, LocalDiskStore
+"""
+
+from atdata.index import (
+    Index,
+    LocalDatasetEntry,
+    BasicIndexEntry,
+    SchemaNamespace,
+    SchemaFieldType,
+    SchemaField,
+    LocalSchemaRecord,
+    REDIS_KEY_DATASET_ENTRY,
+    REDIS_KEY_SCHEMA,
+    _ATDATA_URI_PREFIX,
+    _LEGACY_URI_PREFIX,
+    _kind_str_for_sample_type,
+    _schema_ref_from_type,
+    _make_schema_ref,
+    _parse_schema_ref,
+    _increment_patch,
+    _python_type_to_field_type,
+    _build_schema_record,
+)
+from atdata.stores import (
+    LocalDiskStore,
+    S3DataStore,
+    _s3_env,
+    _s3_from_credentials,
+    _create_s3_write_callbacks,
+)
+from atdata.local._repo_legacy import Repo
+
+# Re-export third-party types that were previously importable from the
+# monolithic local.py (tests reference atdata.local.S3FileSystem, etc.)
+from s3fs import S3FileSystem  # noqa: F401 — re-exported for backward compat
+
+__all__ = [
+    # Public API
+    "LocalDiskStore",
+    "Index",
+    "LocalDatasetEntry",
+    "BasicIndexEntry",
+    "S3DataStore",
+    "Repo",
+    "SchemaNamespace",
+    "SchemaFieldType",
+    "SchemaField",
+    "LocalSchemaRecord",
+    "REDIS_KEY_DATASET_ENTRY",
+    "REDIS_KEY_SCHEMA",
+    # Internal helpers (re-exported for backward compatibility)
+    "_ATDATA_URI_PREFIX",
+    "_LEGACY_URI_PREFIX",
+    "_kind_str_for_sample_type",
+    "_schema_ref_from_type",
+    "_make_schema_ref",
+    "_parse_schema_ref",
+    "_increment_patch",
+    "_python_type_to_field_type",
+    "_build_schema_record",
+    "_s3_env",
+    "_s3_from_credentials",
+    "_create_s3_write_callbacks",
+]

atdata/local/_repo_legacy.py

@@ -0,0 +1,218 @@
+"""Deprecated Repo class for legacy S3 repository operations."""
+
+from atdata import Dataset
+
+from atdata.index._entry import LocalDatasetEntry
+from atdata.stores._s3 import _s3_env, _s3_from_credentials, _create_s3_write_callbacks
+
+from pathlib import Path
+from uuid import uuid4
+from tempfile import TemporaryDirectory
+from typing import Any, BinaryIO, TypeVar, cast
+
+from redis import Redis
+import msgpack
+import webdataset as wds
+import warnings
+
+from atdata._protocols import Packable
+
+T = TypeVar("T", bound=Packable)
+
+
+class Repo:
+    """Repository for storing and managing atdata datasets.
+
+    .. deprecated::
+        Use :class:`Index` with :class:`S3DataStore` instead::
+
+            store = S3DataStore(credentials, bucket="my-bucket")
+            index = Index(redis=redis, data_store=store)
+            entry = index.insert_dataset(ds, name="my-dataset")
+
+    Provides storage of datasets in S3-compatible object storage with Redis-based
+    indexing. Datasets are stored as WebDataset tar files with optional metadata.
+
+    Attributes:
+        s3_credentials: S3 credentials dictionary or None.
+        bucket_fs: S3FileSystem instance or None.
+        hive_path: Path within S3 bucket for storing datasets.
+        hive_bucket: Name of the S3 bucket.
+        index: Index instance for tracking datasets.
+    """
+
+    ##
+
+    def __init__(
+        self,
+        s3_credentials: str | Path | dict[str, Any] | None = None,
+        hive_path: str | Path | None = None,
+        redis: Redis | None = None,
+    ) -> None:
+        """Initialize a repository.
+
+        .. deprecated::
+            Use Index with S3DataStore instead.
+
+        Args:
+            s3_credentials: Path to .env file with S3 credentials, or dict with
+                AWS_ENDPOINT, AWS_ACCESS_KEY_ID, and AWS_SECRET_ACCESS_KEY.
+                If None, S3 functionality will be disabled.
+            hive_path: Path within the S3 bucket to store datasets.
+                Required if s3_credentials is provided.
+            redis: Redis connection for indexing. If None, creates a new connection.
+
+        Raises:
+            ValueError: If hive_path is not provided when s3_credentials is set.
+        """
+        warnings.warn(
+            "Repo is deprecated. Use Index with S3DataStore instead:\n"
+            "  store = S3DataStore(credentials, bucket='my-bucket')\n"
+            "  index = Index(redis=redis, data_store=store)\n"
+            "  entry = index.insert_dataset(ds, name='my-dataset')",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+        if s3_credentials is None:
+            self.s3_credentials = None
+        elif isinstance(s3_credentials, dict):
+            self.s3_credentials = s3_credentials
+        else:
+            self.s3_credentials = _s3_env(s3_credentials)
+
+        if self.s3_credentials is None:
+            self.bucket_fs = None
+        else:
+            self.bucket_fs = _s3_from_credentials(self.s3_credentials)
+
+        if self.bucket_fs is not None:
+            if hive_path is None:
+                raise ValueError("Must specify hive path within bucket")
+            self.hive_path = Path(hive_path)
+            self.hive_bucket = self.hive_path.parts[0]
+        else:
+            self.hive_path = None
+            self.hive_bucket = None
+
+        #
+
+        from atdata.index._index import Index
+
+        self.index = Index(redis=redis)
+
+    ##
+
+    def insert(
+        self,
+        ds: Dataset[T],
+        *,
+        name: str,
+        cache_local: bool = False,
+        schema_ref: str | None = None,
+        **kwargs,
+    ) -> tuple[LocalDatasetEntry, Dataset[T]]:
+        """Insert a dataset into the repository.
+
+        Writes the dataset to S3 as WebDataset tar files, stores metadata,
+        and creates an index entry in Redis.
+
+        Args:
+            ds: The dataset to insert.
+            name: Human-readable name for the dataset.
+            cache_local: If True, write to local temporary storage first, then
+                copy to S3. This can be faster for some workloads.
+            schema_ref: Optional schema reference. If None, generates from sample type.
+            **kwargs: Additional arguments passed to wds.ShardWriter.
+
+        Returns:
+            A tuple of (index_entry, new_dataset) where:
+                - index_entry: LocalDatasetEntry for the stored dataset
+                - new_dataset: Dataset object pointing to the stored copy
+
+        Raises:
+            ValueError: If S3 credentials or hive_path are not configured.
+            RuntimeError: If no shards were written.
+        """
+        if self.s3_credentials is None:
+            raise ValueError(
+                "S3 credentials required for insert(). Initialize Repo with s3_credentials."
+            )
+        if self.hive_bucket is None or self.hive_path is None:
+            raise ValueError(
+                "hive_path required for insert(). Initialize Repo with hive_path."
+            )
+
+        new_uuid = str(uuid4())
+
+        hive_fs = _s3_from_credentials(self.s3_credentials)
+
+        # Write metadata
+        metadata_path = (
+            self.hive_path / "metadata" / f"atdata-metadata--{new_uuid}.msgpack"
+        )
+        # Note: S3 doesn't need directories created beforehand - s3fs handles this
+
+        if ds.metadata is not None:
+            # Use s3:// prefix to ensure s3fs treats this as an S3 path
+            with cast(
+                BinaryIO, hive_fs.open(f"s3://{metadata_path.as_posix()}", "wb")
+            ) as f:
+                meta_packed = msgpack.packb(ds.metadata)
+                f.write(cast(bytes, meta_packed))
+
+        # Write data
+        shard_pattern = (self.hive_path / f"atdata--{new_uuid}--%06d.tar").as_posix()
+
+        written_shards: list[str] = []
+        with TemporaryDirectory() as temp_dir:
+            writer_opener, writer_post = _create_s3_write_callbacks(
+                credentials=self.s3_credentials,
+                temp_dir=temp_dir,
+                written_shards=written_shards,
+                fs=hive_fs,
+                cache_local=cache_local,
+                add_s3_prefix=False,
+            )
+
+            with wds.writer.ShardWriter(
+                shard_pattern,
+                opener=writer_opener,
+                post=writer_post,
+                **kwargs,
+            ) as sink:
+                for sample in ds.ordered(batch_size=None):
+                    sink.write(sample.as_wds)
+
+        # Make a new Dataset object for the written dataset copy
+        if len(written_shards) == 0:
+            raise RuntimeError(
+                "Cannot form new dataset entry -- did not write any shards"
+            )
+
+        elif len(written_shards) < 2:
+            new_dataset_url = (
+                self.hive_path / (Path(written_shards[0]).name)
+            ).as_posix()
+
+        else:
+            shard_s3_format = (
+                (self.hive_path / f"atdata--{new_uuid}").as_posix()
+            ) + "--{shard_id}.tar"
+            shard_id_braced = "{" + f"{0:06d}..{len(written_shards) - 1:06d}" + "}"
+            new_dataset_url = shard_s3_format.format(shard_id=shard_id_braced)
+
+        new_dataset = Dataset[ds.sample_type](
+            url=new_dataset_url,
+            metadata_url=metadata_path.as_posix(),
+        )
+
+        # Add to index (use ds._metadata to avoid network requests)
+        new_entry = self.index.add_entry(
+            new_dataset,
+            name=name,
+            schema_ref=schema_ref,
+            metadata=ds._metadata,
+        )
+
+        return new_entry, new_dataset