atdata 0.2.0a1__py3-none-any.whl → 0.2.3b1__py3-none-any.whl

atdata/local.py

@@ -6,11 +6,13 @@ This module provides a local storage backend for atdata datasets using:
 
 The main classes are:
 - Repo: Manages dataset storage in S3 with Redis indexing
-- Index: Redis-backed index for tracking dataset metadata
-- BasicIndexEntry: Index entry representing a stored dataset
+- LocalIndex: Redis-backed index for tracking dataset metadata
+- LocalDatasetEntry: Index entry representing a stored dataset
 
 This is intended for development and small-scale deployment before
-migrating to the full atproto PDS infrastructure.
+migrating to the full atproto PDS infrastructure. The implementation
+uses ATProto-compatible CIDs for content addressing, enabling seamless
+promotion from local storage to the atmosphere (ATProto network).
 """
 
 ##
@@ -20,8 +22,15 @@ from atdata import (
     PackableSample,
     Dataset,
 )
+from atdata._cid import generate_cid
+from atdata._type_utils import (
+    PRIMITIVE_TYPE_MAP,
+    unwrap_optional,
+    is_ndarray_type,
+    extract_ndarray_dtype,
+)
+from atdata._protocols import AbstractDataStore, Packable
 
-import os
 from pathlib import Path
 from uuid import uuid4
 from tempfile import TemporaryDirectory
@@ -38,142 +47,706 @@ import webdataset as wds
 
 from dataclasses import (
     dataclass,
-    asdict,
     field,
 )
 from typing import (
     Any,
-    Optional,
-    Dict,
     Type,
     TypeVar,
     Generator,
+    Iterator,
     BinaryIO,
+    Optional,
+    Literal,
     cast,
+    get_type_hints,
+    get_origin,
+    get_args,
 )
+from dataclasses import fields, is_dataclass
+from datetime import datetime, timezone
+import json
+import warnings
+
+T = TypeVar("T", bound=PackableSample)
+
+# Redis key prefixes for index entries and schemas
+REDIS_KEY_DATASET_ENTRY = "LocalDatasetEntry"
+REDIS_KEY_SCHEMA = "LocalSchema"
+
 
-T = TypeVar( 'T', bound = PackableSample )
+class SchemaNamespace:
+    """Namespace for accessing loaded schema types as attributes.
+
+    This class provides a module-like interface for accessing dynamically
+    loaded schema types. After calling ``index.load_schema(uri)``, the
+    schema's class becomes available as an attribute on this namespace.
+
+    Examples:
+        >>> index.load_schema("atdata://local/sampleSchema/MySample@1.0.0")
+        >>> MyType = index.types.MySample
+        >>> sample = MyType(field1="hello", field2=42)
+
+    The namespace supports:
+    - Attribute access: ``index.types.MySample``
+    - Iteration: ``for name in index.types: ...``
+    - Length: ``len(index.types)``
+    - Contains check: ``"MySample" in index.types``
+
+    Note:
+        For full IDE autocomplete support, import from the generated module::
+
+            # After load_schema with auto_stubs=True
+            from local.MySample_1_0_0 import MySample
+            sample = MySample(name="hello", value=42)  # IDE knows signature!
+
+        Add ``index.stub_dir`` to your IDE's extraPaths for imports to resolve.
+    """
+
+    def __init__(self) -> None:
+        self._types: dict[str, Type[Packable]] = {}
+
+    def _register(self, name: str, cls: Type[Packable]) -> None:
+        """Register a schema type in the namespace."""
+        self._types[name] = cls
+
+    def __getattr__(self, name: str) -> Any:
+        # Returns Any to avoid IDE complaints about unknown attributes.
+        # For full IDE support, import from the generated module instead.
+        if name.startswith("_"):
+            raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'")
+        if name not in self._types:
+            raise AttributeError(
+                f"Schema '{name}' not loaded. "
+                f"Call index.load_schema() first to load the schema."
+            )
+        return self._types[name]
+
+    def __dir__(self) -> list[str]:
+        return list(self._types.keys()) + ["_types", "_register", "get"]
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._types)
+
+    def __len__(self) -> int:
+        return len(self._types)
+
+    def __contains__(self, name: str) -> bool:
+        return name in self._types
+
+    def __repr__(self) -> str:
+        if not self._types:
+            return "SchemaNamespace(empty)"
+        names = ", ".join(sorted(self._types.keys()))
+        return f"SchemaNamespace({names})"
+
+    def get(self, name: str, default: T | None = None) -> Type[Packable] | T | None:
+        """Get a type by name, returning default if not found.
+
+        Args:
+            name: The schema class name to look up.
+            default: Value to return if not found (default: None).
+
+        Returns:
+            The schema class, or default if not loaded.
+        """
+        return self._types.get(name, default)
+
+
+##
+# Schema types
+
+
+@dataclass
+class SchemaFieldType:
+    """Schema field type definition for local storage.
+
+    Represents a type in the schema type system, supporting primitives,
+    ndarrays, arrays, and references to other schemas.
+    """
+
+    kind: Literal["primitive", "ndarray", "ref", "array"]
+    """The category of type."""
+
+    primitive: Optional[str] = None
+    """For kind='primitive': one of 'str', 'int', 'float', 'bool', 'bytes'."""
+
+    dtype: Optional[str] = None
+    """For kind='ndarray': numpy dtype string (e.g., 'float32')."""
+
+    ref: Optional[str] = None
+    """For kind='ref': URI of referenced schema."""
+
+    items: Optional["SchemaFieldType"] = None
+    """For kind='array': type of array elements."""
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "SchemaFieldType":
+        """Create from a dictionary (e.g., from Redis storage)."""
+        type_str = data.get("$type", "")
+        if "#" in type_str:
+            kind = type_str.split("#")[-1]
+        else:
+            kind = data.get("kind", "primitive")
+
+        items = None
+        if "items" in data and data["items"]:
+            items = cls.from_dict(data["items"])
+
+        return cls(
+            kind=kind,  # type: ignore[arg-type]
+            primitive=data.get("primitive"),
+            dtype=data.get("dtype"),
+            ref=data.get("ref"),
+            items=items,
+        )
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for storage."""
+        result: dict[str, Any] = {"$type": f"local#{self.kind}"}
+        if self.kind == "primitive":
+            result["primitive"] = self.primitive
+        elif self.kind == "ndarray":
+            result["dtype"] = self.dtype
+        elif self.kind == "ref":
+            result["ref"] = self.ref
+        elif self.kind == "array" and self.items:
+            result["items"] = self.items.to_dict()
+        return result
+
+
+@dataclass
+class SchemaField:
+    """Schema field definition for local storage."""
+
+    name: str
+    """Field name."""
+
+    field_type: SchemaFieldType
+    """Type of this field."""
+
+    optional: bool = False
+    """Whether this field can be None."""
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "SchemaField":
+        """Create from a dictionary."""
+        return cls(
+            name=data["name"],
+            field_type=SchemaFieldType.from_dict(data["fieldType"]),
+            optional=data.get("optional", False),
+        )
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for storage."""
+        return {
+            "name": self.name,
+            "fieldType": self.field_type.to_dict(),
+            "optional": self.optional,
+        }
+
+    def __getitem__(self, key: str) -> Any:
+        """Dict-style access for backwards compatibility."""
+        if key == "name":
+            return self.name
+        elif key == "fieldType":
+            return self.field_type.to_dict()
+        elif key == "optional":
+            return self.optional
+        raise KeyError(key)
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Dict-style get() for backwards compatibility."""
+        try:
+            return self[key]
+        except KeyError:
+            return default
+
+
+@dataclass
+class LocalSchemaRecord:
+    """Schema record for local storage.
+
+    Represents a PackableSample schema stored in the local index.
+    Aligns with the atmosphere SchemaRecord structure for seamless promotion.
+    """
+
+    name: str
+    """Schema name (typically the class name)."""
+
+    version: str
+    """Semantic version string (e.g., '1.0.0')."""
+
+    fields: list[SchemaField]
+    """List of field definitions."""
+
+    ref: str
+    """Schema reference URI (atdata://local/sampleSchema/{name}@{version})."""
+
+    description: Optional[str] = None
+    """Human-readable description."""
+
+    created_at: Optional[datetime] = None
+    """When this schema was published."""
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "LocalSchemaRecord":
+        """Create from a dictionary (e.g., from Redis storage)."""
+        created_at = None
+        if "createdAt" in data:
+            try:
+                created_at = datetime.fromisoformat(data["createdAt"])
+            except (ValueError, TypeError):
+                created_at = None  # Invalid datetime format, leave as None
+
+        return cls(
+            name=data["name"],
+            version=data["version"],
+            fields=[SchemaField.from_dict(f) for f in data.get("fields", [])],
+            ref=data.get("$ref", ""),
+            description=data.get("description"),
+            created_at=created_at,
+        )
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for storage."""
+        result: dict[str, Any] = {
+            "name": self.name,
+            "version": self.version,
+            "fields": [f.to_dict() for f in self.fields],
+            "$ref": self.ref,
+        }
+        if self.description:
+            result["description"] = self.description
+        if self.created_at:
+            result["createdAt"] = self.created_at.isoformat()
+        return result
+
+    def __getitem__(self, key: str) -> Any:
+        """Dict-style access for backwards compatibility."""
+        if key == "name":
+            return self.name
+        elif key == "version":
+            return self.version
+        elif key == "fields":
+            return self.fields  # Returns list of SchemaField (also subscriptable)
+        elif key == "$ref":
+            return self.ref
+        elif key == "description":
+            return self.description
+        elif key == "createdAt":
+            return self.created_at.isoformat() if self.created_at else None
+        raise KeyError(key)
+
+    def __contains__(self, key: str) -> bool:
+        """Support 'in' operator for backwards compatibility."""
+        return key in ("name", "version", "fields", "$ref", "description", "createdAt")
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Dict-style get() for backwards compatibility."""
+        try:
+            return self[key]
+        except KeyError:
+            return default
 
 
 ##
 # Helpers
 
-def _kind_str_for_sample_type( st: Type[PackableSample] ) -> str:
-    """Convert a sample type to a fully-qualified string identifier.
+
+def _kind_str_for_sample_type(st: Type[Packable]) -> str:
+    """Return fully-qualified 'module.name' string for a sample type."""
+    return f"{st.__module__}.{st.__name__}"
+
+
+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:
-        st: The sample type class.
+        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:
-        A string in the format 'module.name' identifying the sample type.
+        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:
+        assert fs is not None, "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
+
+
+##
+# Schema helpers
+
+# URI scheme prefixes
+_ATDATA_URI_PREFIX = "atdata://local/sampleSchema/"
+_LEGACY_URI_PREFIX = "local://schemas/"
+
+
+def _schema_ref_from_type(sample_type: Type[Packable], version: str) -> str:
+    """Generate 'atdata://local/sampleSchema/{name}@{version}' reference."""
+    return _make_schema_ref(sample_type.__name__, version)
+
+
+def _make_schema_ref(name: str, version: str) -> str:
+    """Generate schema reference URI from name and version."""
+    return f"{_ATDATA_URI_PREFIX}{name}@{version}"
+
+
+def _parse_schema_ref(ref: str) -> tuple[str, str]:
+    """Parse schema reference into (name, version).
+
+    Supports both new format: 'atdata://local/sampleSchema/{name}@{version}'
+    and legacy format: 'local://schemas/{module.Class}@{version}'
     """
-    return f'{st.__module__}.{st.__name__}'
+    if ref.startswith(_ATDATA_URI_PREFIX):
+        path = ref[len(_ATDATA_URI_PREFIX) :]
+    elif ref.startswith(_LEGACY_URI_PREFIX):
+        path = ref[len(_LEGACY_URI_PREFIX) :]
+    else:
+        raise ValueError(f"Invalid schema reference: {ref}")
+
+    if "@" not in path:
+        raise ValueError(f"Schema reference must include version (@version): {ref}")
+
+    name, version = path.rsplit("@", 1)
+    # For legacy format, extract just the class name from module.Class
+    if "." in name:
+        name = name.rsplit(".", 1)[1]
+    return name, version
+
+
+def _parse_semver(version: str) -> tuple[int, int, int]:
+    """Parse semantic version string into (major, minor, patch) tuple."""
+    parts = version.split(".")
+    if len(parts) != 3:
+        raise ValueError(f"Invalid semver format: {version}")
+    return int(parts[0]), int(parts[1]), int(parts[2])
+
+
+def _increment_patch(version: str) -> str:
+    """Increment patch version: 1.0.0 -> 1.0.1"""
+    major, minor, patch = _parse_semver(version)
+    return f"{major}.{minor}.{patch + 1}"
+
+
+def _python_type_to_field_type(python_type: Any) -> dict:
+    """Convert Python type annotation to schema field type dict."""
+    if python_type in PRIMITIVE_TYPE_MAP:
+        return {
+            "$type": "local#primitive",
+            "primitive": PRIMITIVE_TYPE_MAP[python_type],
+        }
+
+    if is_ndarray_type(python_type):
+        return {"$type": "local#ndarray", "dtype": extract_ndarray_dtype(python_type)}
+
+    origin = get_origin(python_type)
+    if origin is list:
+        args = get_args(python_type)
+        items = (
+            _python_type_to_field_type(args[0])
+            if args
+            else {"$type": "local#primitive", "primitive": "str"}
+        )
+        return {"$type": "local#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."
+        )
 
-def _decode_bytes_dict( d: dict[bytes, bytes] ) -> dict[str, str]:
-    """Decode a dictionary with byte keys and values to strings.
+    raise TypeError(f"Unsupported type for schema field: {python_type}")
 
-    Redis returns dictionaries with bytes keys/values, this converts them to strings.
+
+def _build_schema_record(
+    sample_type: Type[Packable],
+    *,
+    version: str,
+    description: str | None = None,
+) -> dict:
+    """Build a schema record dict from a PackableSample type.
 
     Args:
-        d: Dictionary with bytes keys and values.
+        sample_type: The PackableSample subclass to introspect.
+        version: Semantic version string.
+        description: Optional human-readable description. If None, uses the
+            class docstring.
 
     Returns:
-        Dictionary with UTF-8 decoded string keys and values.
+        Schema record dict suitable for Redis storage.
+
+    Raises:
+        ValueError: If sample_type is not a dataclass.
+        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)")
+
+    # Use docstring as fallback for description
+    if description is None:
+        description = sample_type.__doc__
+
+    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_type, is_optional = unwrap_optional(field_type)
+        field_type_dict = _python_type_to_field_type(field_type)
+
+        field_defs.append(
+            {
+                "name": f.name,
+                "fieldType": field_type_dict,
+                "optional": is_optional,
+            }
+        )
+
     return {
-        k.decode('utf-8'): v.decode('utf-8')
-        for k, v in d.items()
+        "name": sample_type.__name__,
+        "version": version,
+        "fields": field_defs,
+        "description": description,
+        "createdAt": datetime.now(timezone.utc).isoformat(),
     }
 
 
 ##
 # Redis object model
 
+
 @dataclass
-class BasicIndexEntry:
-    """Index entry for a dataset stored in the repository.
+class LocalDatasetEntry:
+    """Index entry for a dataset stored in the local repository.
+
+    Implements the IndexEntry protocol for compatibility with AbstractIndex.
+    Uses dual identity: a content-addressable CID (ATProto-compatible) and
+    a human-readable name.
+
+    The CID is generated from the entry's content (schema_ref + data_urls),
+    ensuring the same data produces the same CID whether stored locally or
+    in the atmosphere. This enables seamless promotion from local to ATProto.
 
-    Tracks metadata about a dataset stored in S3, including its location,
-    type, and unique identifier.
+    Attributes:
+        name: Human-readable name for this dataset.
+        schema_ref: Reference to the schema for this dataset.
+        data_urls: WebDataset URLs for the data.
+        metadata: Arbitrary metadata dictionary, or None if not set.
     """
+
     ##
 
-    wds_url: str
-    """WebDataset URL for the dataset tar files, for use with atdata.Dataset."""
+    name: str
+    """Human-readable name for this dataset."""
+
+    schema_ref: str
+    """Reference to the schema for this dataset."""
 
-    sample_kind: str
-    """Fully-qualified sample type name (e.g., 'module.ClassName')."""
+    data_urls: list[str]
+    """WebDataset URLs for the data."""
 
-    metadata_url: str | None
-    """S3 URL to the dataset's metadata msgpack file, if any."""
+    metadata: dict | None = None
+    """Arbitrary metadata dictionary, or None if not set."""
+
+    _cid: str | None = field(default=None, repr=False)
+    """Content identifier (ATProto-compatible CID). Generated from content if not provided."""
+
+    # Legacy field for backwards compatibility during migration
+    _legacy_uuid: str | None = field(default=None, repr=False)
+    """Legacy UUID for backwards compatibility with existing Redis entries."""
+
+    def __post_init__(self):
+        """Generate CID from content if not provided."""
+        if self._cid is None:
+            self._cid = self._generate_cid()
+
+    def _generate_cid(self) -> str:
+        """Generate ATProto-compatible CID from entry content."""
+        # CID is based on schema_ref and data_urls - the identity of the dataset
+        content = {
+            "schema_ref": self.schema_ref,
+            "data_urls": self.data_urls,
+        }
+        return generate_cid(content)
+
+    @property
+    def cid(self) -> str:
+        """Content identifier (ATProto-compatible CID)."""
+        assert self._cid is not None
+        return self._cid
 
-    uuid: str = field( default_factory = lambda: str( uuid4() ) )
-    """Unique identifier for this dataset entry. Defaults to a new UUID if not provided."""
+    # Legacy compatibility
 
-    def write_to( self, redis: Redis ):
+    @property
+    def wds_url(self) -> str:
+        """Legacy property: returns first data URL for backwards compatibility."""
+        return self.data_urls[0] if self.data_urls else ""
+
+    @property
+    def sample_kind(self) -> str:
+        """Legacy property: returns schema_ref for backwards compatibility."""
+        return self.schema_ref
+
+    def write_to(self, redis: Redis):
         """Persist this index entry to Redis.
 
-        Stores the entry as a Redis hash with key 'BasicIndexEntry:{uuid}'.
+        Stores the entry as a Redis hash with key '{REDIS_KEY_DATASET_ENTRY}:{cid}'.
 
         Args:
             redis: Redis connection to write to.
         """
-        save_key = f'BasicIndexEntry:{self.uuid}'
-        # Filter out None values - Redis doesn't accept None
-        data = {k: v for k, v in asdict(self).items() if v is not None}
-        # redis-py typing uses untyped dict, so type checker complains about dict[str, Any]
-        redis.hset( save_key, mapping = data )  # type: ignore[arg-type]
+        save_key = f"{REDIS_KEY_DATASET_ENTRY}:{self.cid}"
+        data = {
+            "name": self.name,
+            "schema_ref": self.schema_ref,
+            "data_urls": msgpack.packb(self.data_urls),  # Serialize list
+            "cid": self.cid,
+        }
+        if self.metadata is not None:
+            data["metadata"] = msgpack.packb(self.metadata)
+        if self._legacy_uuid is not None:
+            data["legacy_uuid"] = self._legacy_uuid
+
+        redis.hset(save_key, mapping=data)  # type: ignore[arg-type]
+
+    @classmethod
+    def from_redis(cls, redis: Redis, cid: str) -> "LocalDatasetEntry":
+        """Load an entry from Redis by CID.
+
+        Args:
+            redis: Redis connection to read from.
+            cid: Content identifier of the entry to load.
 
-def _s3_env( credentials_path: str | Path ) -> dict[str, Any]:
-    """Load S3 credentials from a .env file.
+        Returns:
+            LocalDatasetEntry loaded from Redis.
+
+        Raises:
+            KeyError: If entry not found.
+        """
+        save_key = f"{REDIS_KEY_DATASET_ENTRY}:{cid}"
+        raw_data = redis.hgetall(save_key)
+        if not raw_data:
+            raise KeyError(f"{REDIS_KEY_DATASET_ENTRY} not found: {cid}")
+
+        # Decode string fields, keep binary fields as bytes for msgpack
+        raw_data_typed = cast(dict[bytes, bytes], raw_data)
+        name = raw_data_typed[b"name"].decode("utf-8")
+        schema_ref = raw_data_typed[b"schema_ref"].decode("utf-8")
+        cid_value = raw_data_typed.get(b"cid", b"").decode("utf-8") or None
+        legacy_uuid = raw_data_typed.get(b"legacy_uuid", b"").decode("utf-8") or None
+
+        # Deserialize msgpack fields (stored as raw bytes)
+        data_urls = msgpack.unpackb(raw_data_typed[b"data_urls"])
+        metadata = None
+        if b"metadata" in raw_data_typed:
+            metadata = msgpack.unpackb(raw_data_typed[b"metadata"])
+
+        return cls(
+            name=name,
+            schema_ref=schema_ref,
+            data_urls=data_urls,
+            metadata=metadata,
+            _cid=cid_value,
+            _legacy_uuid=legacy_uuid,
+        )
+
+
+# Backwards compatibility alias
+BasicIndexEntry = LocalDatasetEntry
+
+
+def _s3_env(credentials_path: str | Path) -> dict[str, Any]:
+    """Load S3 credentials from .env file.
 
     Args:
-        credentials_path: Path to .env file containing S3 credentials.
+        credentials_path: Path to .env file containing AWS_ENDPOINT,
+            AWS_ACCESS_KEY_ID, and AWS_SECRET_ACCESS_KEY.
 
     Returns:
-        Dictionary with AWS_ENDPOINT, AWS_ACCESS_KEY_ID, and AWS_SECRET_ACCESS_KEY.
+        Dict with the three required credential keys.
 
     Raises:
-        AssertionError: If required credentials are missing from the file.
+        ValueError: If any required key is missing from the .env file.
     """
-    ##
-    credentials_path = Path( credentials_path )
-    env_values = dotenv_values( credentials_path )
-    assert 'AWS_ENDPOINT' in env_values
-    assert 'AWS_ACCESS_KEY_ID' in env_values
-    assert 'AWS_SECRET_ACCESS_KEY' in env_values
-
-    return {
-        k: env_values[k]
-        for k in (
-            'AWS_ENDPOINT',
-            'AWS_ACCESS_KEY_ID',
-            'AWS_SECRET_ACCESS_KEY',
+    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)}"
         )
-    }
 
-def _s3_from_credentials( creds: str | Path | dict ) -> S3FileSystem:
-    """Create an S3FileSystem from credentials.
+    return {k: env_values[k] for k in required_keys}
 
-    Args:
-        creds: Either a path to a .env file with credentials, or a dict
-            containing AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, and optionally
-            AWS_ENDPOINT.
 
-    Returns:
-        Configured S3FileSystem instance.
-    """
-    ##
-    if not isinstance( creds, dict ):
-        creds = _s3_env( creds )
+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']
+        "key": creds["AWS_ACCESS_KEY_ID"],
+        "secret": creds["AWS_SECRET_ACCESS_KEY"],
     }
-    if 'AWS_ENDPOINT' in creds:
-        kwargs['endpoint_url'] = creds['AWS_ENDPOINT']
+    if "AWS_ENDPOINT" in creds:
+        kwargs["endpoint_url"] = creds["AWS_ENDPOINT"]
 
     return S3FileSystem(**kwargs)
 
@@ -181,9 +754,17 @@ def _s3_from_credentials( creds: str | Path | dict ) -> S3FileSystem:
 ##
 # Classes
 
+
 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.
 
@@ -197,17 +778,17 @@ class Repo:
 
     ##
 
-    def __init__( self,
-                #
-                s3_credentials: str | Path | dict[str, Any] | None = None,
-                hive_path: str | Path | None = None,
-                redis: Redis | None = None,
-                #
-                #
-                **kwargs
-            ) -> None:
+    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.
@@ -215,45 +796,55 @@ class Repo:
             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.
-            **kwargs: Additional arguments (reserved for future use).
 
         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 ):
+        elif isinstance(s3_credentials, dict):
             self.s3_credentials = s3_credentials
         else:
-            self.s3_credentials = _s3_env( s3_credentials )
+            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 )
+            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 )
+                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
-        
+
         #
 
-        self.index = Index( redis = redis )
+        self.index = Index(redis=redis)
 
     ##
 
-    def insert( self, ds: Dataset[T],
-               #
-               cache_local: bool = False,
-               #
-                **kwargs
-            ) -> tuple[BasicIndexEntry, Dataset[T]]:
+    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,
@@ -261,130 +852,102 @@ class Repo:
 
         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: BasicIndexEntry for the stored dataset
+                - index_entry: LocalDatasetEntry for the stored dataset
                 - new_dataset: Dataset object pointing to the stored copy
 
         Raises:
-            AssertionError: If S3 credentials or hive_path are not configured.
+            ValueError: If S3 credentials or hive_path are not configured.
             RuntimeError: If no shards were written.
         """
-        
-        assert self.s3_credentials is not None
-        assert self.hive_bucket is not None
-        assert self.hive_path is not None
+        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() )
+        new_uuid = str(uuid4())
 
-        hive_fs = _s3_from_credentials( self.s3_credentials )
+        hive_fs = _s3_from_credentials(self.s3_credentials)
 
         # Write metadata
         metadata_path = (
-            self.hive_path
-            / 'metadata'
-            / f'atdata-metadata--{new_uuid}.msgpack'
+            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 )
+            with cast(
+                BinaryIO, hive_fs.open(f"s3://{metadata_path.as_posix()}", "wb")
+            ) as f:
+                meta_packed = msgpack.packb(ds.metadata)
                 assert meta_packed is not None
-                f.write( cast( bytes, meta_packed ) )
-
+                f.write(cast(bytes, meta_packed))
 
         # Write data
-        shard_pattern = (
-            self.hive_path
-            / f'atdata--{new_uuid}--%06d.tar'
-        ).as_posix()
+        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,
+            )
 
-            if cache_local:
-                # For cache_local, we need to use boto3 directly to avoid s3fs async issues with moto
-                import boto3
-
-                # Create boto3 client from credentials
-                s3_client_kwargs = {
-                    'aws_access_key_id': self.s3_credentials['AWS_ACCESS_KEY_ID'],
-                    'aws_secret_access_key': self.s3_credentials['AWS_SECRET_ACCESS_KEY']
-                }
-                if 'AWS_ENDPOINT' in self.s3_credentials:
-                    s3_client_kwargs['endpoint_url'] = self.s3_credentials['AWS_ENDPOINT']
-                s3_client = boto3.client('s3', **s3_client_kwargs)
-
-                def _writer_opener( p: str ):
-                    local_cache_path = Path( temp_dir ) / p
-                    local_cache_path.parent.mkdir( parents = True, exist_ok = True )
-                    return open( local_cache_path, 'wb' )
-                writer_opener = _writer_opener
-
-                def _writer_post( p: str ):
-                    local_cache_path = Path( temp_dir ) / p
-
-                    # Copy to S3 using boto3 client (avoids s3fs async issues)
-                    path_parts = Path( p ).parts
-                    bucket = path_parts[0]
-                    key = str( Path( *path_parts[1:] ) )
-
-                    with open( local_cache_path, 'rb' ) as f_in:
-                        s3_client.put_object( Bucket=bucket, Key=key, Body=f_in.read() )
-
-                    # Delete local cache file
-                    local_cache_path.unlink()
-
-                    written_shards.append( p )
-                writer_post = _writer_post
-
-            else:
-                # Use s3:// prefix to ensure s3fs treats paths as S3 paths
-                writer_opener = lambda s: cast( BinaryIO, hive_fs.open( f's3://{s}', 'wb' ) )
-                writer_post = lambda s: written_shards.append( s )
-
-            written_shards = []
             with wds.writer.ShardWriter(
                 shard_pattern,
-                opener = writer_opener,
-                post = writer_post,
+                opener=writer_opener,
+                post=writer_post,
                 **kwargs,
             ) as sink:
-                for sample in ds.ordered( batch_size = None ):
-                    sink.write( sample.as_wds )
+                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:
+        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 )
+                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 )
+                (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(),
+            url=new_dataset_url,
+            metadata_url=metadata_path.as_posix(),
         )
 
-        # Add to index
-        new_entry = self.index.add_entry( new_dataset, uuid = new_uuid )
+        # 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
 
@@ -392,24 +955,43 @@ class Repo:
 class Index:
     """Redis-backed index for tracking datasets in a repository.
 
-    Maintains a registry of BasicIndexEntry objects in Redis, allowing
-    enumeration and lookup of stored datasets.
+    Implements the AbstractIndex protocol. Maintains a registry of
+    LocalDatasetEntry objects in Redis, allowing enumeration and lookup
+    of stored datasets.
+
+    When initialized with a data_store, insert_dataset() will write dataset
+    shards to storage before indexing. Without a data_store, insert_dataset()
+    only indexes existing URLs.
 
     Attributes:
         _redis: Redis connection for index storage.
+        _data_store: Optional AbstractDataStore for writing dataset shards.
     """
 
     ##
 
-    def __init__( self,
-                redis: Redis | None = None,
-                **kwargs
-            ) -> None:
+    def __init__(
+        self,
+        redis: Redis | None = None,
+        data_store: AbstractDataStore | None = None,
+        auto_stubs: bool = False,
+        stub_dir: Path | str | None = None,
+        **kwargs,
+    ) -> None:
         """Initialize an index.
 
         Args:
             redis: Redis connection to use. If None, creates a new connection
                 using the provided kwargs.
+            data_store: Optional data store for writing dataset shards.
+                If provided, insert_dataset() will write shards to this store.
+                If None, insert_dataset() only indexes existing URLs.
+            auto_stubs: If True, automatically generate .pyi stub files when
+                schemas are accessed via get_schema() or decode_schema().
+                This enables IDE autocomplete for dynamically decoded types.
+            stub_dir: Directory to write stub files. Only used if auto_stubs
+                is True or if this parameter is provided (which implies auto_stubs).
+                Defaults to ~/.atdata/stubs/ if not specified.
             **kwargs: Additional arguments passed to Redis() constructor if
                 redis is None.
         """
@@ -418,75 +1000,721 @@ class Index:
         if redis is not None:
             self._redis = redis
         else:
-            self._redis: Redis = Redis( **kwargs )
+            self._redis: Redis = Redis(**kwargs)
+
+        self._data_store = data_store
+
+        # Initialize stub manager if auto-stubs enabled
+        # Providing stub_dir implies auto_stubs=True
+        if auto_stubs or stub_dir is not None:
+            from ._stub_manager import StubManager
+
+            self._stub_manager: StubManager | None = StubManager(stub_dir=stub_dir)
+        else:
+            self._stub_manager = None
+
+        # Initialize schema namespace for load_schema/schemas API
+        self._schema_namespace = SchemaNamespace()
+
+    @property
+    def data_store(self) -> AbstractDataStore | None:
+        """The data store for writing shards, or None if index-only."""
+        return self._data_store
 
     @property
-    def all_entries( self ) -> list[BasicIndexEntry]:
-        """Get all index entries as a list.
+    def stub_dir(self) -> Path | None:
+        """Directory where stub files are written, or None if auto-stubs disabled.
+
+        Use this path to configure your IDE for type checking support:
+        - VS Code/Pylance: Add to python.analysis.extraPaths in settings.json
+        - PyCharm: Mark as Sources Root
+        - mypy: Add to mypy_path in mypy.ini
+        """
+        if self._stub_manager is not None:
+            return self._stub_manager.stub_dir
+        return None
+
+    @property
+    def types(self) -> SchemaNamespace:
+        """Namespace for accessing loaded schema types.
+
+        After calling :meth:`load_schema`, schema types become available
+        as attributes on this namespace.
+
+        Examples:
+            >>> index.load_schema("atdata://local/sampleSchema/MySample@1.0.0")
+            >>> MyType = index.types.MySample
+            >>> sample = MyType(name="hello", value=42)
+
+        Returns:
+            SchemaNamespace containing all loaded schema types.
+        """
+        return self._schema_namespace
+
+    def load_schema(self, ref: str) -> Type[Packable]:
+        """Load a schema and make it available in the types namespace.
+
+        This method decodes the schema, optionally generates a Python module
+        for IDE support (if auto_stubs is enabled), and registers the type
+        in the :attr:`types` namespace for easy access.
+
+        Args:
+            ref: Schema reference string (atdata://local/sampleSchema/... or
+                legacy local://schemas/...).
 
         Returns:
-            List of all BasicIndexEntry objects in the index.
+            The decoded PackableSample subclass. Also available via
+            ``index.types.<ClassName>`` after this call.
+
+        Raises:
+            KeyError: If schema not found.
+            ValueError: If schema cannot be decoded.
+
+        Examples:
+            >>> # Load and use immediately
+            >>> MyType = index.load_schema("atdata://local/sampleSchema/MySample@1.0.0")
+            >>> sample = MyType(name="hello", value=42)
+            >>>
+            >>> # Or access later via namespace
+            >>> index.load_schema("atdata://local/sampleSchema/OtherType@1.0.0")
+            >>> other = index.types.OtherType(data="test")
         """
-        return list( self.entries )
+        # Decode the schema (uses generated module if auto_stubs enabled)
+        cls = self.decode_schema(ref)
+
+        # Register in namespace using the class name
+        self._schema_namespace._register(cls.__name__, cls)
+
+        return cls
+
+    def get_import_path(self, ref: str) -> str | None:
+        """Get the import path for a schema's generated module.
+
+        When auto_stubs is enabled, this returns the import path that can
+        be used to import the schema type with full IDE support.
+
+        Args:
+            ref: Schema reference string.
+
+        Returns:
+            Import path like "local.MySample_1_0_0", or None if auto_stubs
+            is disabled.
+
+        Examples:
+            >>> index = LocalIndex(auto_stubs=True)
+            >>> ref = index.publish_schema(MySample, version="1.0.0")
+            >>> index.load_schema(ref)
+            >>> print(index.get_import_path(ref))
+            local.MySample_1_0_0
+            >>> # Then in your code:
+            >>> # from local.MySample_1_0_0 import MySample
+        """
+        if self._stub_manager is None:
+            return None
+
+        from ._stub_manager import _extract_authority
+
+        name, version = _parse_schema_ref(ref)
+        schema_dict = self.get_schema(ref)
+        authority = _extract_authority(schema_dict.get("$ref"))
+
+        safe_version = version.replace(".", "_")
+        module_name = f"{name}_{safe_version}"
+
+        return f"{authority}.{module_name}"
+
+    def list_entries(self) -> list[LocalDatasetEntry]:
+        """Get all index entries as a materialized list.
+
+        Returns:
+            List of all LocalDatasetEntry objects in the index.
+        """
+        return list(self.entries)
+
+    # Legacy alias for backwards compatibility
+    @property
+    def all_entries(self) -> list[LocalDatasetEntry]:
+        """Get all index entries as a list (deprecated, use list_entries())."""
+        return self.list_entries()
 
     @property
-    def entries( self ) -> Generator[BasicIndexEntry, None, None]:
+    def entries(self) -> Generator[LocalDatasetEntry, None, None]:
         """Iterate over all index entries.
 
-        Scans Redis for all BasicIndexEntry keys and yields them one at a time.
+        Scans Redis for LocalDatasetEntry keys and yields them one at a time.
 
         Yields:
-            BasicIndexEntry objects from the index.
+            LocalDatasetEntry objects from the index.
         """
-        ##
-        for key in self._redis.scan_iter( match = 'BasicIndexEntry:*' ):
-            # hgetall returns dict[bytes, bytes] which we decode to dict[str, str]
-            cur_entry_data = _decode_bytes_dict( cast(dict[bytes, bytes], self._redis.hgetall( key )) )
-            
-            # Provide default None for optional fields that may be missing
-            # Type checker complains about None in dict[str, str], but BasicIndexEntry accepts it
-            cur_entry_data: dict[str, Any] = dict( **cur_entry_data )
-            cur_entry_data.setdefault('metadata_url', None)
-            
-            cur_entry = BasicIndexEntry( **cur_entry_data )
-            yield cur_entry
-
-        return
-
-    def add_entry( self, ds: Dataset,
-                uuid: str | None = None,
-            ) -> BasicIndexEntry:
+        prefix = f"{REDIS_KEY_DATASET_ENTRY}:"
+        for key in self._redis.scan_iter(match=f"{prefix}*"):
+            key_str = key.decode("utf-8") if isinstance(key, bytes) else key
+            cid = key_str[len(prefix) :]
+            yield LocalDatasetEntry.from_redis(self._redis, cid)
+
+    def add_entry(
+        self,
+        ds: Dataset,
+        *,
+        name: str,
+        schema_ref: str | None = None,
+        metadata: dict | None = None,
+    ) -> LocalDatasetEntry:
         """Add a dataset to the index.
 
-        Creates a BasicIndexEntry for the dataset and persists it to Redis.
+        Creates a LocalDatasetEntry for the dataset and persists it to Redis.
 
         Args:
             ds: The dataset to add to the index.
-            uuid: Optional UUID for the entry. If None, a new UUID is generated.
+            name: Human-readable name for the dataset.
+            schema_ref: Optional schema reference. If None, generates from sample type.
+            metadata: Optional metadata dictionary. If None, uses ds._metadata if available.
 
         Returns:
-            The created BasicIndexEntry object.
+            The created LocalDatasetEntry object.
         """
         ##
-        temp_sample_kind = _kind_str_for_sample_type( ds.sample_type )
+        if schema_ref is None:
+            schema_ref = (
+                f"local://schemas/{_kind_str_for_sample_type(ds.sample_type)}@1.0.0"
+            )
+
+        # Normalize URL to list
+        data_urls = [ds.url]
+
+        # Use provided metadata, or fall back to dataset's cached metadata
+        # (avoid triggering network requests via ds.metadata property)
+        entry_metadata = metadata if metadata is not None else ds._metadata
+
+        entry = LocalDatasetEntry(
+            name=name,
+            schema_ref=schema_ref,
+            data_urls=data_urls,
+            metadata=entry_metadata,
+        )
+
+        entry.write_to(self._redis)
+
+        return entry
+
+    def get_entry(self, cid: str) -> LocalDatasetEntry:
+        """Get an entry by its CID.
+
+        Args:
+            cid: Content identifier of the entry.
 
-        if uuid is None:
-            ret_data = BasicIndexEntry(
-                wds_url = ds.url,
-                sample_kind = temp_sample_kind,
-                metadata_url = ds.metadata_url,
+        Returns:
+            LocalDatasetEntry for the given CID.
+
+        Raises:
+            KeyError: If entry not found.
+        """
+        return LocalDatasetEntry.from_redis(self._redis, cid)
+
+    def get_entry_by_name(self, name: str) -> LocalDatasetEntry:
+        """Get an entry by its human-readable name.
+
+        Args:
+            name: Human-readable name of the entry.
+
+        Returns:
+            LocalDatasetEntry with the given name.
+
+        Raises:
+            KeyError: If no entry with that name exists.
+        """
+        for entry in self.entries:
+            if entry.name == name:
+                return entry
+        raise KeyError(f"No entry with name: {name}")
+
+    # AbstractIndex protocol methods
+
+    def insert_dataset(
+        self,
+        ds: Dataset,
+        *,
+        name: str,
+        schema_ref: str | None = None,
+        **kwargs,
+    ) -> LocalDatasetEntry:
+        """Insert a dataset into the index (AbstractIndex protocol).
+
+        If a data_store was provided at initialization, writes dataset shards
+        to storage first, then indexes the new URLs. Otherwise, indexes the
+        dataset's existing URL.
+
+        Args:
+            ds: The Dataset to register.
+            name: Human-readable name for the dataset.
+            schema_ref: Optional schema reference.
+            **kwargs: Additional options:
+                - metadata: Optional metadata dict
+                - prefix: Storage prefix (default: dataset name)
+                - cache_local: If True, cache writes locally first
+
+        Returns:
+            IndexEntry for the inserted dataset.
+        """
+        metadata = kwargs.get("metadata")
+
+        if self._data_store is not None:
+            # Write shards to data store, then index the new URLs
+            prefix = kwargs.get("prefix", name)
+            cache_local = kwargs.get("cache_local", False)
+
+            written_urls = self._data_store.write_shards(
+                ds,
+                prefix=prefix,
+                cache_local=cache_local,
             )
+
+            # Generate schema_ref if not provided
+            if schema_ref is None:
+                schema_ref = _schema_ref_from_type(ds.sample_type, version="1.0.0")
+
+            # Create entry with the written URLs
+            entry_metadata = metadata if metadata is not None else ds._metadata
+            entry = LocalDatasetEntry(
+                name=name,
+                schema_ref=schema_ref,
+                data_urls=written_urls,
+                metadata=entry_metadata,
+            )
+            entry.write_to(self._redis)
+            return entry
+
+        # No data store - just index the existing URL
+        return self.add_entry(ds, name=name, schema_ref=schema_ref, metadata=metadata)
+
+    def get_dataset(self, ref: str) -> LocalDatasetEntry:
+        """Get a dataset entry by name (AbstractIndex protocol).
+
+        Args:
+            ref: Dataset name.
+
+        Returns:
+            IndexEntry for the dataset.
+
+        Raises:
+            KeyError: If dataset not found.
+        """
+        return self.get_entry_by_name(ref)
+
+    @property
+    def datasets(self) -> Generator[LocalDatasetEntry, None, None]:
+        """Lazily iterate over all dataset entries (AbstractIndex protocol).
+
+        Yields:
+            IndexEntry for each dataset.
+        """
+        return self.entries
+
+    def list_datasets(self) -> list[LocalDatasetEntry]:
+        """Get all dataset entries as a materialized list (AbstractIndex protocol).
+
+        Returns:
+            List of IndexEntry for each dataset.
+        """
+        return self.list_entries()
+
+    # Schema operations
+
+    def _get_latest_schema_version(self, name: str) -> str | None:
+        """Get the latest version for a schema by name, or None if not found."""
+        latest_version: tuple[int, int, int] | None = None
+        latest_version_str: str | None = None
+
+        prefix = f"{REDIS_KEY_SCHEMA}:"
+        for key in self._redis.scan_iter(match=f"{prefix}*"):
+            key_str = key.decode("utf-8") if isinstance(key, bytes) else key
+            schema_id = key_str[len(prefix) :]
+
+            if "@" not in schema_id:
+                continue
+
+            schema_name, version_str = schema_id.rsplit("@", 1)
+            # Handle legacy format: module.Class -> Class
+            if "." in schema_name:
+                schema_name = schema_name.rsplit(".", 1)[1]
+
+            if schema_name != name:
+                continue
+
+            try:
+                version_tuple = _parse_semver(version_str)
+                if latest_version is None or version_tuple > latest_version:
+                    latest_version = version_tuple
+                    latest_version_str = version_str
+            except ValueError:
+                continue
+
+        return latest_version_str
+
+    def publish_schema(
+        self,
+        sample_type: type,
+        *,
+        version: str | None = None,
+        description: str | None = None,
+    ) -> str:
+        """Publish a schema for a sample type to Redis.
+
+        Args:
+            sample_type: A Packable type (@packable-decorated or PackableSample subclass).
+            version: Semantic version string (e.g., '1.0.0'). If None,
+                auto-increments from the latest published version (patch bump),
+                or starts at '1.0.0' if no previous version exists.
+            description: Optional human-readable description. If None, uses
+                the class docstring.
+
+        Returns:
+            Schema reference string: 'atdata://local/sampleSchema/{name}@{version}'.
+
+        Raises:
+            ValueError: If sample_type is not a dataclass.
+            TypeError: If sample_type doesn't satisfy the Packable protocol,
+                or if a field type is not supported.
+        """
+        # Validate that sample_type satisfies Packable protocol at runtime
+        # This catches non-packable types early with a clear error message
+        try:
+            # Check protocol compliance by verifying required methods exist
+            if not (
+                hasattr(sample_type, "from_data")
+                and hasattr(sample_type, "from_bytes")
+                and callable(getattr(sample_type, "from_data", None))
+                and callable(getattr(sample_type, "from_bytes", None))
+            ):
+                raise TypeError(
+                    f"{sample_type.__name__} does not satisfy the Packable protocol. "
+                    "Use @packable decorator or inherit from PackableSample."
+                )
+        except AttributeError:
+            raise TypeError(
+                f"sample_type must be a class, got {type(sample_type).__name__}"
+            )
+
+        # Auto-increment version if not specified
+        if version is None:
+            latest = self._get_latest_schema_version(sample_type.__name__)
+            if latest is None:
+                version = "1.0.0"
+            else:
+                version = _increment_patch(latest)
+
+        schema_record = _build_schema_record(
+            sample_type,
+            version=version,
+            description=description,
+        )
+
+        schema_ref = _schema_ref_from_type(sample_type, version)
+        name, _ = _parse_schema_ref(schema_ref)
+
+        # Store in Redis
+        redis_key = f"{REDIS_KEY_SCHEMA}:{name}@{version}"
+        schema_json = json.dumps(schema_record)
+        self._redis.set(redis_key, schema_json)
+
+        return schema_ref
+
+    def get_schema(self, ref: str) -> dict:
+        """Get a schema record by reference (AbstractIndex protocol).
+
+        Args:
+            ref: Schema reference string. Supports both new format
+                (atdata://local/sampleSchema/{name}@{version}) and legacy
+                format (local://schemas/{module.Class}@{version}).
+
+        Returns:
+            Schema record as a dictionary with keys 'name', 'version',
+            'fields', '$ref', etc.
+
+        Raises:
+            KeyError: If schema not found.
+            ValueError: If reference format is invalid.
+        """
+        name, version = _parse_schema_ref(ref)
+        redis_key = f"{REDIS_KEY_SCHEMA}:{name}@{version}"
+
+        schema_json = self._redis.get(redis_key)
+        if schema_json is None:
+            raise KeyError(f"Schema not found: {ref}")
+
+        if isinstance(schema_json, bytes):
+            schema_json = schema_json.decode("utf-8")
+
+        schema = json.loads(schema_json)
+        schema["$ref"] = _make_schema_ref(name, version)
+
+        # Auto-generate stub if enabled
+        if self._stub_manager is not None:
+            record = LocalSchemaRecord.from_dict(schema)
+            self._stub_manager.ensure_stub(record)
+
+        return schema
+
+    def get_schema_record(self, ref: str) -> LocalSchemaRecord:
+        """Get a schema record as LocalSchemaRecord object.
+
+        Use this when you need the full LocalSchemaRecord with typed properties.
+        For Protocol-compliant dict access, use get_schema() instead.
+
+        Args:
+            ref: Schema reference string.
+
+        Returns:
+            LocalSchemaRecord with schema details.
+
+        Raises:
+            KeyError: If schema not found.
+            ValueError: If reference format is invalid.
+        """
+        schema = self.get_schema(ref)
+        return LocalSchemaRecord.from_dict(schema)
+
+    @property
+    def schemas(self) -> Generator[LocalSchemaRecord, None, None]:
+        """Iterate over all schema records in this index.
+
+        Yields:
+            LocalSchemaRecord for each schema.
+        """
+        prefix = f"{REDIS_KEY_SCHEMA}:"
+        for key in self._redis.scan_iter(match=f"{prefix}*"):
+            key_str = key.decode("utf-8") if isinstance(key, bytes) else key
+            # Extract name@version from key
+            schema_id = key_str[len(prefix) :]
+
+            schema_json = self._redis.get(key)
+            if schema_json is None:
+                continue
+
+            if isinstance(schema_json, bytes):
+                schema_json = schema_json.decode("utf-8")
+
+            schema = json.loads(schema_json)
+            # Handle legacy keys that have module.Class format
+            if "." in schema_id.split("@")[0]:
+                name = schema_id.split("@")[0].rsplit(".", 1)[1]
+                version = schema_id.split("@")[1]
+                schema["$ref"] = _make_schema_ref(name, version)
+            else:
+                # schema_id is already "name@version"
+                name, version = schema_id.rsplit("@", 1)
+                schema["$ref"] = _make_schema_ref(name, version)
+            yield LocalSchemaRecord.from_dict(schema)
+
+    def list_schemas(self) -> list[dict]:
+        """Get all schema records as a materialized list (AbstractIndex protocol).
+
+        Returns:
+            List of schema records as dictionaries.
+        """
+        return [record.to_dict() for record in self.schemas]
+
+    def decode_schema(self, ref: str) -> Type[Packable]:
+        """Reconstruct a Python PackableSample type from a stored schema.
+
+        This method enables loading datasets without knowing the sample type
+        ahead of time. The index retrieves the schema record and dynamically
+        generates a PackableSample subclass matching the schema definition.
+
+        If auto_stubs is enabled, a Python module will be generated and the
+        class will be imported from it, providing full IDE autocomplete support.
+        The returned class has proper type information that IDEs can understand.
+
+        Args:
+            ref: Schema reference string (atdata://local/sampleSchema/... or
+                legacy local://schemas/...).
+
+        Returns:
+            A PackableSample subclass - either imported from a generated module
+            (if auto_stubs is enabled) or dynamically created.
+
+        Raises:
+            KeyError: If schema not found.
+            ValueError: If schema cannot be decoded.
+        """
+        schema_dict = self.get_schema(ref)
+
+        # If auto_stubs is enabled, generate module and import class from it
+        if self._stub_manager is not None:
+            cls = self._stub_manager.ensure_module(schema_dict)
+            if cls is not None:
+                return cls
+
+        # Fall back to dynamic type generation
+        from atdata._schema_codec import schema_to_type
+
+        return schema_to_type(schema_dict)
+
+    def decode_schema_as(self, ref: str, type_hint: type[T]) -> type[T]:
+        """Decode a schema with explicit type hint for IDE support.
+
+        This is a typed wrapper around decode_schema() that preserves the
+        type information for IDE autocomplete. Use this when you have a
+        stub file for the schema and want full IDE support.
+
+        Args:
+            ref: Schema reference string.
+            type_hint: The stub type to use for type hints. Import this from
+                the generated stub file.
+
+        Returns:
+            The decoded type, cast to match the type_hint for IDE support.
+
+        Examples:
+            >>> # After enabling auto_stubs and configuring IDE extraPaths:
+            >>> from local.MySample_1_0_0 import MySample
+            >>>
+            >>> # This gives full IDE autocomplete:
+            >>> DecodedType = index.decode_schema_as(ref, MySample)
+            >>> sample = DecodedType(text="hello", value=42)  # IDE knows signature!
+
+        Note:
+            The type_hint is only used for static type checking - at runtime,
+            the actual decoded type from the schema is returned. Ensure the
+            stub matches the schema to avoid runtime surprises.
+        """
+        from typing import cast
+
+        return cast(type[T], self.decode_schema(ref))
+
+    def clear_stubs(self) -> int:
+        """Remove all auto-generated stub files.
+
+        Only works if auto_stubs was enabled when creating the Index.
+
+        Returns:
+            Number of stub files removed, or 0 if auto_stubs is disabled.
+        """
+        if self._stub_manager is not None:
+            return self._stub_manager.clear_stubs()
+        return 0
+
+
+# Backwards compatibility alias
+LocalIndex = Index
+
+
+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:
-            ret_data = BasicIndexEntry(
-                wds_url = ds.url,
-                sample_kind = temp_sample_kind,
-                metadata_url = ds.metadata_url,
-                uuid = uuid,
+            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,
+        **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.
+            **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] = []
+
+        with TemporaryDirectory() as temp_dir:
+            writer_opener, writer_post = _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,
             )
 
-        ret_data.write_to( self._redis )
+            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)
+
+        if len(written_shards) == 0:
+            raise RuntimeError("No shards written")
+
+        return written_shards
 
-        return ret_data
+    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
 
 
-#
+#