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

atdata/{local → index}/_index.py

@@ -7,8 +7,8 @@ from atdata import (
 )
 from atdata._protocols import AbstractDataStore, Packable
 
-from atdata.local._entry import LocalDatasetEntry
-from atdata.local._schema import (
+from atdata.index._entry import LocalDatasetEntry
+from atdata.index._schema import (
     SchemaNamespace,
     LocalSchemaRecord,
     _schema_ref_from_type,
@@ -21,6 +21,7 @@ from atdata.local._schema import (
 from pathlib import Path
 from typing import (
     Any,
+    Iterable,
     Type,
     TypeVar,
     Generator,
@@ -41,8 +42,8 @@ class Index:
     """Unified index for tracking datasets across multiple repositories.
 
     Implements the AbstractIndex protocol. Maintains a registry of
-    dataset entries across a built-in ``"local"`` repository, optional
-    named repositories, and an optional atmosphere (ATProto) backend.
+    dataset entries across named repositories (always including a built-in
+    ``"local"`` repository) and an optional atmosphere (ATProto) backend.
 
     The ``"local"`` repository is always present and uses the storage backend
     determined by the ``provider`` argument. When no provider is given, defaults
@@ -52,14 +53,12 @@ class Index:
     Additional named repositories can be mounted via the ``repos`` parameter,
     each pairing an IndexProvider with an optional data store.
 
-    An AtmosphereClient is available by default for anonymous read-only
+    An Atmosphere is available by default for anonymous read-only
     resolution of ``@handle/dataset`` paths. Pass an authenticated client
     for write operations, or ``atmosphere=None`` to disable.
 
     Attributes:
-        _provider: IndexProvider for the built-in ``"local"`` repository.
-        _data_store: Optional AbstractDataStore for the local repository.
-        _repos: Named repositories beyond ``"local"``.
+        _repos: All repositories keyed by name. ``"local"`` is always present.
         _atmosphere: Optional atmosphere backend for ATProto operations.
     """
 
@@ -105,7 +104,7 @@ class Index:
             atmosphere: ATProto client for distributed network operations.
                 - Default (sentinel): creates an anonymous read-only client
                   lazily on first access.
-                - ``AtmosphereClient`` instance: uses that client directly.
+                - ``Atmosphere`` instance: uses that client directly.
                 - ``None``: disables atmosphere backend entirely.
             auto_stubs: If True, automatically generate .pyi stub files when
                 schemas are accessed via get_schema() or decode_schema().
@@ -146,12 +145,13 @@ class Index:
         ##
 
         from atdata.providers._base import IndexProvider as _IP
+        from atdata.repository import Repository as _Repo
 
+        # Resolve the local provider
         if isinstance(provider, str):
-            # String-based provider selection
             from atdata.providers._factory import create_provider
 
-            self._provider: _IP = create_provider(
+            local_provider: _IP = create_provider(
                 provider, path=path, dsn=dsn, redis=redis, **kwargs
             )
         elif provider is not None:
@@ -160,27 +160,24 @@ class Index:
                     f"provider must be an IndexProvider or backend name string, "
                     f"got {type(provider).__name__}"
                 )
-            self._provider = provider
+            local_provider = provider
         elif redis is not None:
-            # Explicit Redis connection provided
             from atdata.providers._redis import RedisProvider
 
-            self._provider = RedisProvider(redis)
+            local_provider = RedisProvider(redis)
         elif kwargs:
-            # kwargs provided — assume Redis constructor args for compat
             from atdata.providers._redis import RedisProvider
 
-            self._provider = RedisProvider(Redis(**kwargs))
+            local_provider = RedisProvider(Redis(**kwargs))
         else:
-            # Default: zero-dependency SQLite
             from atdata.providers._sqlite import SqliteProvider
 
-            self._provider = SqliteProvider()
-
-        self._data_store = data_store
+            local_provider = SqliteProvider()
 
-        # Validate and store named repositories
-        from atdata.repository import Repository as _Repo
+        # Build the unified repos dict with "local" always present
+        self._repos: dict[str, _Repo] = {
+            "local": _Repo(provider=local_provider, data_store=data_store),
+        }
 
         if repos is not None:
             if "local" in repos:
@@ -194,9 +191,7 @@ class Index:
                         f"repos[{name!r}] must be a Repository, "
                         f"got {type(repo).__name__}"
                     )
-            self._repos: dict[str, _Repo] = dict(repos)
-        else:
-            self._repos = {}
+            self._repos.update(repos)
 
         # Atmosphere backend (lazy or explicit)
         from atdata.repository import _AtmosphereBackend
@@ -230,10 +225,10 @@ class Index:
         """Get the atmosphere backend, lazily creating anonymous client if needed."""
         if self._atmosphere_deferred and self._atmosphere is None:
             try:
-                from atdata.atmosphere.client import AtmosphereClient
+                from atdata.atmosphere.client import Atmosphere
                 from atdata.repository import _AtmosphereBackend
 
-                client = AtmosphereClient()
+                client = Atmosphere()
                 self._atmosphere = _AtmosphereBackend(client)
             except ImportError:
                 # atproto package not installed -- atmosphere unavailable
@@ -289,13 +284,13 @@ class Index:
         return ("local", ref, None)
 
     @property
-    def repos(self) -> dict[str, Repository]:
-        """Named repositories mounted on this index (excluding ``"local"``)."""
+    def repos(self) -> dict[str, "Repository"]:
+        """All repositories mounted on this index (including ``"local"``)."""
         return dict(self._repos)
 
     @property
     def atmosphere(self) -> Any:
-        """The AtmosphereClient for this index, or None if disabled.
+        """The Atmosphere for this index, or None if disabled.
 
         Returns the underlying client (not the internal backend wrapper).
         """
@@ -304,10 +299,15 @@ class Index:
             return backend.client
         return None
 
+    @property
+    def _provider(self) -> "IndexProvider":  # noqa: F821
+        """IndexProvider for the ``"local"`` repository (backward compat)."""
+        return self._repos["local"].provider
+
     @property
     def provider(self) -> "IndexProvider":  # noqa: F821
-        """The storage provider backing this index."""
-        return self._provider
+        """The storage provider backing the ``"local"`` repository."""
+        return self._repos["local"].provider
 
     @property
     def _redis(self) -> Redis:
@@ -318,17 +318,23 @@ class Index:
         """
         from atdata.providers._redis import RedisProvider
 
-        if isinstance(self._provider, RedisProvider):
-            return self._provider.redis
+        prov = self._repos["local"].provider
+        if isinstance(prov, RedisProvider):
+            return prov.redis
         raise AttributeError(
             "Index._redis is only available with a Redis provider. "
             "Use index.provider instead."
         )
 
+    @property
+    def _data_store(self) -> AbstractDataStore | None:
+        """Data store for the ``"local"`` repository (backward compat)."""
+        return self._repos["local"].data_store
+
     @property
     def data_store(self) -> AbstractDataStore | None:
         """The data store for writing shards, or None if index-only."""
-        return self._data_store
+        return self._repos["local"].data_store
 
     @property
     def stub_dir(self) -> Path | None:
@@ -351,7 +357,7 @@ class Index:
         as attributes on this namespace.
 
         Examples:
-            >>> index.load_schema("atdata://local/sampleSchema/MySample@1.0.0")
+            >>> index.load_schema("atdata://local/schema/MySample@1.0.0")
             >>> MyType = index.types.MySample
             >>> sample = MyType(name="hello", value=42)
 
@@ -368,7 +374,7 @@ class Index:
         in the :attr:`types` namespace for easy access.
 
         Args:
-            ref: Schema reference string (atdata://local/sampleSchema/... or
+            ref: Schema reference string (atdata://local/schema/... or
                 legacy local://schemas/...).
 
         Returns:
@@ -381,11 +387,11 @@ class Index:
 
         Examples:
             >>> # Load and use immediately
-            >>> MyType = index.load_schema("atdata://local/sampleSchema/MySample@1.0.0")
+            >>> MyType = index.load_schema("atdata://local/schema/MySample@1.0.0")
             >>> sample = MyType(field1="hello", field2=42)
             >>>
             >>> # Or access later via namespace
-            >>> index.load_schema("atdata://local/sampleSchema/OtherType@1.0.0")
+            >>> index.load_schema("atdata://local/schema/OtherType@1.0.0")
             >>> other = index.types.OtherType(data="test")
         """
         # Decode the schema (uses generated module if auto_stubs enabled)
@@ -513,6 +519,23 @@ class Index:
 
     # AbstractIndex protocol methods
 
+    @staticmethod
+    def _ensure_schema_stored(
+        schema_ref: str,
+        sample_type: type,
+        provider: "IndexProvider",  # noqa: F821
+    ) -> None:
+        """Persist the schema definition if not already stored.
+
+        Called during dataset insertion so that ``decode_schema()`` can
+        reconstruct the type later without the caller needing to publish
+        the schema separately.
+        """
+        schema_name, version = _parse_schema_ref(schema_ref)
+        if provider.get_schema_json(schema_name, version) is None:
+            record = _build_schema_record(sample_type, version=version)
+            provider.store_schema(schema_name, version, json.dumps(record))
+
     def _insert_dataset_to_provider(
         self,
         ds: Dataset,
@@ -543,6 +566,8 @@ class Index:
             if schema_ref is None:
                 schema_ref = _schema_ref_from_type(ds.sample_type, version="1.0.0")
 
+            self._ensure_schema_stored(schema_ref, ds.sample_type, provider)
+
             entry_metadata = metadata if metadata is not None else ds._metadata
             entry = LocalDatasetEntry(
                 name=name,
@@ -557,6 +582,8 @@ class Index:
         if schema_ref is None:
             schema_ref = _schema_ref_from_type(ds.sample_type, version="1.0.0")
 
+        self._ensure_schema_stored(schema_ref, ds.sample_type, provider)
+
         data_urls = [ds.url]
         entry_metadata = metadata if metadata is not None else ds._metadata
 
@@ -612,17 +639,6 @@ class Index:
                 ds, name=resolved_name, schema_ref=schema_ref, **kwargs
             )
 
-        if backend_key == "local":
-            return self._insert_dataset_to_provider(
-                ds,
-                name=resolved_name,
-                schema_ref=schema_ref,
-                provider=self._provider,
-                store=self._data_store,
-                **kwargs,
-            )
-
-        # Named repository
         repo = self._repos.get(backend_key)
         if repo is None:
             raise KeyError(f"Unknown repository {backend_key!r} in name {name!r}")
@@ -635,6 +651,117 @@ class Index:
             **kwargs,
         )
 
+    def write(
+        self,
+        samples: Iterable,
+        *,
+        name: str,
+        schema_ref: str | None = None,
+        description: str | None = None,
+        tags: list[str] | None = None,
+        license: str | None = None,
+        maxcount: int = 10_000,
+        maxsize: int | None = None,
+        metadata: dict | None = None,
+        manifest: bool = False,
+    ) -> "IndexEntry":
+        """Write samples and create an index entry in one step.
+
+        This is the primary method for publishing data. It serializes
+        samples to WebDataset tar files, stores them via the appropriate
+        backend, and creates an index entry.
+
+        The target backend is determined by the *name* prefix:
+
+        - Bare name (e.g., ``"mnist"``): writes to the local repository.
+        - ``"@handle/name"``: writes and publishes to the atmosphere.
+        - ``"repo/name"``: writes to a named repository.
+
+        When the local backend has no ``data_store`` configured, a
+        ``LocalDiskStore`` is created automatically at
+        ``~/.atdata/data/`` so that samples have persistent storage.
+
+        .. note::
+
+            This method is synchronous. Samples are written to a temporary
+            location first, then copied to permanent storage by the backend.
+            Avoid passing lazily-evaluated iterators that depend on external
+            state that may change during the call.
+
+        Args:
+            samples: Iterable of ``Packable`` samples. Must be non-empty.
+            name: Dataset name, optionally prefixed with target.
+            schema_ref: Optional schema reference. Auto-generated if ``None``.
+            description: Optional dataset description (atmosphere only).
+            tags: Optional tags for discovery (atmosphere only).
+            license: Optional license identifier (atmosphere only).
+            maxcount: Max samples per shard. Default: 10,000.
+            maxsize: Max bytes per shard. Default: ``None``.
+            metadata: Optional metadata dict stored with the entry.
+            manifest: If True, write per-shard manifest sidecar files
+                alongside each tar. Default: ``False``.
+
+        Returns:
+            IndexEntry for the created dataset.
+
+        Raises:
+            ValueError: If *samples* is empty.
+
+        Examples:
+            >>> index = Index()
+            >>> samples = [MySample(key="0", text="hello")]
+            >>> entry = index.write(samples, name="my-dataset")
+        """
+        import tempfile
+
+        from atdata.dataset import write_samples
+
+        backend_key, resolved_name, _ = self._resolve_prefix(name)
+
+        # Resolve the target repo's data store; auto-create LocalDiskStore
+        # for repos that have no store so write() always persists data.
+        repo = self._repos.get(backend_key)
+        effective_store = repo.data_store if repo is not None else None
+        needs_auto_store = repo is not None and effective_store is None
+
+        if needs_auto_store and backend_key != "_atmosphere":
+            from atdata.stores._disk import LocalDiskStore
+
+            effective_store = LocalDiskStore()
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            tmp_path = Path(tmp_dir) / "data.tar"
+            ds = write_samples(
+                samples,
+                tmp_path,
+                maxcount=maxcount,
+                maxsize=maxsize,
+                manifest=manifest,
+            )
+
+            # When we auto-created a store, write directly through it
+            # rather than via insert_dataset (which would just index
+            # the temp path).
+            if needs_auto_store and repo is not None:
+                return self._insert_dataset_to_provider(
+                    ds,
+                    name=resolved_name,
+                    schema_ref=schema_ref,
+                    provider=repo.provider,
+                    store=effective_store,
+                    metadata=metadata,
+                )
+
+            return self.insert_dataset(
+                ds,
+                name=name,
+                schema_ref=schema_ref,
+                metadata=metadata,
+                description=description,
+                tags=tags,
+                license=license,
+            )
+
     def get_dataset(self, ref: str) -> "IndexEntry":
         """Get a dataset entry by name or prefixed reference.
 
@@ -659,14 +786,10 @@ class Index:
             if atmo is None:
                 raise ValueError(
                     f"Atmosphere backend required for path {ref!r} but not available. "
-                    "Install 'atproto' or pass an AtmosphereClient."
+                    "Install 'atproto' or pass an Atmosphere."
                 )
             return atmo.get_dataset(resolved_ref)
 
-        if backend_key == "local":
-            return self._provider.get_entry_by_name(resolved_ref)
-
-        # Named repository
         repo = self._repos.get(backend_key)
         if repo is None:
             raise KeyError(f"Unknown repository {backend_key!r} in ref {ref!r}")
@@ -676,14 +799,13 @@ class Index:
     def datasets(self) -> Generator["IndexEntry", None, None]:
         """Lazily iterate over all dataset entries across local repositories.
 
-        Yields entries from the ``"local"`` repository and all named
-        repositories. Atmosphere entries are not included (use
+        Yields entries from all mounted repositories (``"local"`` and named).
+        Atmosphere entries are not included (use
         ``list_datasets(repo="_atmosphere")`` for those).
 
         Yields:
             IndexEntry for each dataset.
         """
-        yield from self._provider.iter_entries()
         for repo in self._repos.values():
             yield from repo.provider.iter_entries()
 
@@ -702,9 +824,6 @@ class Index:
         if repo is None:
             return list(self.datasets)
 
-        if repo == "local":
-            return self.list_entries()
-
         if repo == "_atmosphere":
             atmo = self._get_atmosphere()
             if atmo is None:
@@ -740,7 +859,7 @@ class Index:
                 the class docstring.
 
         Returns:
-            Schema reference string: 'atdata://local/sampleSchema/{name}@{version}'.
+            Schema reference string: 'atdata://local/schema/{name}@{version}'.
 
         Raises:
             ValueError: If sample_type is not a dataclass.
@@ -794,7 +913,7 @@ class Index:
 
         Args:
             ref: Schema reference string. Supports both new format
-                (atdata://local/sampleSchema/{name}@{version}) and legacy
+                (atdata://local/schema/{name}@{version}) and legacy
                 format (local://schemas/{module.Class}@{version}).
 
         Returns:
@@ -871,7 +990,7 @@ class Index:
         The returned class has proper type information that IDEs can understand.
 
         Args:
-            ref: Schema reference string (atdata://local/sampleSchema/... or
+            ref: Schema reference string (atdata://local/schema/... or
                 legacy local://schemas/...).
 
         Returns:
@@ -938,3 +1057,142 @@ class Index:
         if self._stub_manager is not None:
             return self._stub_manager.clear_stubs()
         return 0
+
+    # -- Atmosphere promotion --
+
+    def promote_entry(
+        self,
+        entry_name: str,
+        *,
+        name: str | None = None,
+        description: str | None = None,
+        tags: list[str] | None = None,
+        license: str | None = None,
+    ) -> str:
+        """Promote a locally-indexed dataset to the atmosphere.
+
+        Looks up the entry by name in the local index, resolves its
+        schema, and publishes both schema and dataset record to ATProto
+        via the index's atmosphere backend.
+
+        Args:
+            entry_name: Name of the local dataset entry to promote.
+            name: Override name for the atmosphere record. Defaults to
+                the local entry name.
+            description: Optional description for the dataset.
+            tags: Optional tags for discovery.
+            license: Optional license identifier.
+
+        Returns:
+            AT URI of the created atmosphere dataset record.
+
+        Raises:
+            ValueError: If atmosphere backend is not available, or
+                the local entry has no data URLs.
+            KeyError: If the entry or its schema is not found.
+
+        Examples:
+            >>> index = Index(atmosphere=client)
+            >>> uri = index.promote_entry("mnist-train")
+        """
+        from atdata.promote import _find_or_publish_schema
+        from atdata.atmosphere import DatasetPublisher
+        from atdata._schema_codec import schema_to_type
+
+        atmo = self._get_atmosphere()
+        if atmo is None:
+            raise ValueError("Atmosphere backend required but not available.")
+
+        entry = self.get_entry_by_name(entry_name)
+        if not entry.data_urls:
+            raise ValueError(f"Local entry {entry_name!r} has no data URLs")
+
+        schema_record = self.get_schema(entry.schema_ref)
+        sample_type = schema_to_type(schema_record)
+        schema_version = schema_record.get("version", "1.0.0")
+
+        atmosphere_schema_uri = _find_or_publish_schema(
+            sample_type,
+            schema_version,
+            atmo.client,
+            description=schema_record.get("description"),
+        )
+
+        publisher = DatasetPublisher(atmo.client)
+        uri = publisher.publish_with_urls(
+            urls=entry.data_urls,
+            schema_uri=atmosphere_schema_uri,
+            name=name or entry.name,
+            description=description,
+            tags=tags,
+            license=license,
+            metadata=entry.metadata,
+        )
+        return str(uri)
+
+    def promote_dataset(
+        self,
+        dataset: Dataset,
+        *,
+        name: str,
+        sample_type: type | None = None,
+        schema_version: str = "1.0.0",
+        description: str | None = None,
+        tags: list[str] | None = None,
+        license: str | None = None,
+    ) -> str:
+        """Publish a Dataset directly to the atmosphere.
+
+        Publishes the schema (with deduplication) and creates a dataset
+        record on ATProto. Uses the index's atmosphere backend.
+
+        Args:
+            dataset: The Dataset to publish.
+            name: Name for the atmosphere dataset record.
+            sample_type: Sample type for schema publishing. Inferred from
+                ``dataset.sample_type`` if not provided.
+            schema_version: Semantic version for the schema. Default: ``"1.0.0"``.
+            description: Optional description for the dataset.
+            tags: Optional tags for discovery.
+            license: Optional license identifier.
+
+        Returns:
+            AT URI of the created atmosphere dataset record.
+
+        Raises:
+            ValueError: If atmosphere backend is not available.
+
+        Examples:
+            >>> index = Index(atmosphere=client)
+            >>> ds = atdata.load_dataset("./data.tar", MySample, split="train")
+            >>> uri = index.promote_dataset(ds, name="my-dataset")
+        """
+        from atdata.promote import _find_or_publish_schema
+        from atdata.atmosphere import DatasetPublisher
+
+        atmo = self._get_atmosphere()
+        if atmo is None:
+            raise ValueError("Atmosphere backend required but not available.")
+
+        st = sample_type or dataset.sample_type
+
+        atmosphere_schema_uri = _find_or_publish_schema(
+            st,
+            schema_version,
+            atmo.client,
+            description=description,
+        )
+
+        data_urls = dataset.list_shards()
+
+        publisher = DatasetPublisher(atmo.client)
+        uri = publisher.publish_with_urls(
+            urls=data_urls,
+            schema_uri=atmosphere_schema_uri,
+            name=name,
+            description=description,
+            tags=tags,
+            license=license,
+            metadata=dataset._metadata,
+        )
+        return str(uri)

atdata/{local → index}/_schema.py

@@ -26,7 +26,7 @@ from typing import (
 T = TypeVar("T", bound=Packable)
 
 # URI scheme prefixes
-_ATDATA_URI_PREFIX = "atdata://local/sampleSchema/"
+_ATDATA_URI_PREFIX = "atdata://local/schema/"
 _LEGACY_URI_PREFIX = "local://schemas/"
 
 
@@ -37,7 +37,7 @@ class SchemaNamespace:
     Supports attribute access, iteration, ``len()``, and ``in`` checks.
 
     Examples:
-        >>> index.load_schema("atdata://local/sampleSchema/MySample@1.0.0")
+        >>> index.load_schema("atdata://local/schema/MySample@1.0.0")
         >>> MyType = index.types.MySample
         >>> sample = MyType(field1="hello", field2=42)
 
@@ -207,7 +207,7 @@ class LocalSchemaRecord:
     """List of field definitions."""
 
     ref: str
-    """Schema reference URI (atdata://local/sampleSchema/{name}@{version})."""
+    """Schema reference URI (atdata://local/schema/{name}@{version})."""
 
     description: Optional[str] = None
     """Human-readable description."""
@@ -259,7 +259,7 @@ def _kind_str_for_sample_type(st: Type[Packable]) -> str:
 
 
 def _schema_ref_from_type(sample_type: Type[Packable], version: str) -> str:
-    """Generate 'atdata://local/sampleSchema/{name}@{version}' reference."""
+    """Generate 'atdata://local/schema/{name}@{version}' reference."""
     return _make_schema_ref(sample_type.__name__, version)
 
 
@@ -271,7 +271,7 @@ def _make_schema_ref(name: str, version: str) -> str:
 def _parse_schema_ref(ref: str) -> tuple[str, str]:
     """Parse schema reference into (name, version).
 
-    Supports both new format: 'atdata://local/sampleSchema/{name}@{version}'
+    Supports both new format: 'atdata://local/schema/{name}@{version}'
     and legacy format: 'local://schemas/{module.Class}@{version}'
     """
     if ref.startswith(_ATDATA_URI_PREFIX):