datachain 0.30.5__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/lib/model_store.py

@@ -1,11 +1,8 @@
 import inspect
-import logging
-from typing import Any, ClassVar, Optional
+from typing import Any, ClassVar
 
 from pydantic import BaseModel
 
-logger = logging.getLogger(__name__)
-
 
 class ModelStore:
     store: ClassVar[dict[str, dict[int, type[BaseModel]]]] = {}
@@ -14,7 +11,7 @@ class ModelStore:
     def get_version(cls, model: type[BaseModel]) -> int:
         if not hasattr(model, "_version"):
             return 0
-        return model._version
+        return model._version  # type: ignore[attr-defined]
 
     @classmethod
     def get_name(cls, model) -> str:
@@ -39,7 +36,7 @@ class ModelStore:
                 cls.register(anno)
 
     @classmethod
-    def get(cls, name: str, version: Optional[int] = None) -> Optional[type]:
+    def get(cls, name: str, version: int | None = None) -> type | None:
         class_dict = cls.store.get(name, None)
         if class_dict is None:
             return None
@@ -77,7 +74,7 @@ class ModelStore:
         )
 
     @staticmethod
-    def to_pydantic(val) -> Optional[type[BaseModel]]:
+    def to_pydantic(val) -> type[BaseModel] | None:
         if val is None or not ModelStore.is_pydantic(val):
             return None
         return val
@@ -98,6 +95,21 @@ class ModelStore:
         (e.g. from by-value cloudpickle in workers) reports built state but
         nested model field schemas aren't fully resolved yet.
         """
+        visited: set[type[BaseModel]] = set()
+        visiting: set[type[BaseModel]] = set()
+
+        def visit(model: type[BaseModel]) -> None:
+            if model in visited or model in visiting:
+                return
+            visiting.add(model)
+            for field in model.model_fields.values():
+                child = cls.to_pydantic(field.annotation)
+                if child is not None:
+                    visit(child)
+            visiting.remove(model)
+            model.model_rebuild(force=True)
+            visited.add(model)
+
         for versions in cls.store.values():
             for model in versions.values():
-                model.model_rebuild(force=True)
+                visit(model)

datachain/lib/namespaces.py

@@ -1,12 +1,14 @@
-from typing import Optional
-
-from datachain.error import NamespaceCreateNotAllowedError
-from datachain.namespace import Namespace
+from datachain.error import (
+    NamespaceCreateNotAllowedError,
+    NamespaceDeleteNotAllowedError,
+)
+from datachain.lib.projects import delete as delete_project
+from datachain.namespace import Namespace, parse_name
 from datachain.query import Session
 
 
 def create(
-    name: str, descr: Optional[str] = None, session: Optional[Session] = None
+    name: str, descr: str | None = None, session: Session | None = None
 ) -> Namespace:
     """
     Creates a new namespace.
@@ -38,7 +40,7 @@ def create(
     return session.catalog.metastore.create_namespace(name, descr)
 
 
-def get(name: str, session: Optional[Session] = None) -> Namespace:
+def get(name: str, session: Session | None = None) -> Namespace:
     """
     Gets a namespace by name.
     If the namespace is not found, a `NamespaceNotFoundError` is raised.
@@ -57,7 +59,7 @@ def get(name: str, session: Optional[Session] = None) -> Namespace:
     return session.catalog.metastore.get_namespace(name)
 
 
-def ls(session: Optional[Session] = None) -> list[Namespace]:
+def ls(session: Session | None = None) -> list[Namespace]:
     """
     Gets a list of all namespaces.
 
@@ -71,3 +73,53 @@ def ls(session: Optional[Session] = None) -> list[Namespace]:
         ```
     """
     return Session.get(session).catalog.metastore.list_namespaces()
+
+
+def delete_namespace(name: str, session: Session | None = None) -> None:
+    """
+    Removes a namespace by name.
+
+    Raises:
+        NamespaceNotFoundError: If the namespace does not exist.
+        NamespaceDeleteNotAllowedError: If the namespace is non-empty,
+            is the default namespace, or is a system namespace,
+            as these cannot be removed.
+
+    Parameters:
+        name: The name of the namespace.
+        session: Session to use for getting project.
+
+    Example:
+        ```py
+        import datachain as dc
+        dc.delete_namespace("dev")
+        ```
+    """
+    session = Session.get(session)
+    metastore = session.catalog.metastore
+
+    namespace_name, project_name = parse_name(name)
+
+    if project_name:
+        return delete_project(project_name, namespace_name, session)
+
+    namespace = metastore.get_namespace(name)
+
+    if name == metastore.system_namespace_name:
+        raise NamespaceDeleteNotAllowedError(
+            f"Namespace {metastore.system_namespace_name} cannot be removed"
+        )
+
+    if name == metastore.default_namespace_name:
+        raise NamespaceDeleteNotAllowedError(
+            f"Namespace {metastore.default_namespace_name} cannot be removed"
+        )
+
+    num_projects = metastore.count_projects(namespace.id)
+    if num_projects > 0:
+        raise NamespaceDeleteNotAllowedError(
+            f"Namespace cannot be removed. It contains {num_projects} project(s). "
+            "Please remove the project(s) first."
+        )
+
+    metastore.remove_namespace(namespace.id)

datachain/lib/projects.py

@@ -1,6 +1,4 @@
-from typing import Optional
-
-from datachain.error import ProjectCreateNotAllowedError
+from datachain.error import ProjectCreateNotAllowedError, ProjectDeleteNotAllowedError
 from datachain.project import Project
 from datachain.query import Session
 
@@ -8,8 +6,8 @@ from datachain.query import Session
 def create(
     namespace: str,
     name: str,
-    descr: Optional[str] = None,
-    session: Optional[Session] = None,
+    descr: str | None = None,
+    session: Session | None = None,
 ) -> Project:
     """
     Creates a new project under a specified namespace.
@@ -42,7 +40,7 @@ def create(
     return session.catalog.metastore.create_project(namespace, name, descr)
 
 
-def get(name: str, namespace: str, session: Optional[Session]) -> Project:
+def get(name: str, namespace: str, session: Session | None) -> Project:
     """
     Gets a project by name in some namespace.
     If the project is not found, a `ProjectNotFoundError` is raised.
@@ -62,9 +60,7 @@ def get(name: str, namespace: str, session: Optional[Session]) -> Project:
     return Session.get(session).catalog.metastore.get_project(name, namespace)
 
 
-def ls(
-    namespace: Optional[str] = None, session: Optional[Session] = None
-) -> list[Project]:
+def ls(namespace: str | None = None, session: Session | None = None) -> list[Project]:
     """
     Gets a list of projects in a specific namespace or from all namespaces.
 
@@ -86,3 +82,49 @@ def ls(
         namespace_id = session.catalog.metastore.get_namespace(namespace).id
 
     return session.catalog.metastore.list_projects(namespace_id)
+
+
+def delete(name: str, namespace: str, session: Session | None = None) -> None:
+    """
+    Removes a project by name within a namespace.
+
+    Raises:
+        ProjectNotFoundError: If the project does not exist.
+        ProjectDeleteNotAllowedError: If the project is non-empty,
+            is the default project, or is a listing project,
+            as these cannot be removed.
+
+    Parameters:
+        name : The name of the project.
+        namespace : The name of the namespace.
+        session : Session to use for getting project.
+
+    Example:
+        ```py
+        import datachain as dc
+        dc.delete_project("my-project", "local")
+        ```
+    """
+    session = Session.get(session)
+    metastore = session.catalog.metastore
+
+    project = metastore.get_project(name, namespace)
+
+    if metastore.is_listing_project(name, namespace):
+        raise ProjectDeleteNotAllowedError(
+            f"Project {metastore.listing_project_name} cannot be removed"
+        )
+
+    if metastore.is_default_project(name, namespace):
+        raise ProjectDeleteNotAllowedError(
+            f"Project {metastore.default_project_name} cannot be removed"
+        )
+
+    num_datasets = metastore.count_datasets(project.id)
+    if num_datasets > 0:
+        raise ProjectDeleteNotAllowedError(
+            f"Project cannot be removed. It contains {num_datasets} dataset(s). "
+            "Please remove the dataset(s) first."
+        )
+
+    metastore.remove_project(project.id)

datachain/lib/pytorch.py

@@ -1,9 +1,9 @@
 import logging
 import os
 import weakref
-from collections.abc import Generator, Iterable, Iterator
+from collections.abc import Callable, Generator, Iterable, Iterator
 from contextlib import closing
-from typing import TYPE_CHECKING, Any, Callable, Optional
+from typing import TYPE_CHECKING, Any
 
 from PIL import Image
 from torch import float32
@@ -43,13 +43,13 @@ class PytorchDataset(IterableDataset):
     def __init__(
         self,
         name: str,
-        version: Optional[str] = None,
-        catalog: Optional["Catalog"] = None,
-        transform: Optional["Transform"] = None,
-        tokenizer: Optional[Callable] = None,
-        tokenizer_kwargs: Optional[dict[str, Any]] = None,
+        version: str | None = None,
+        catalog: Catalog | None = None,
+        transform: "Transform | None" = None,
+        tokenizer: Callable | None = None,
+        tokenizer_kwargs: dict[str, Any] | None = None,
         num_samples: int = 0,
-        dc_settings: Optional[Settings] = None,
+        dc_settings: Settings | None = None,
         remove_prefetched: bool = False,
     ):
         """
@@ -74,6 +74,7 @@ class PytorchDataset(IterableDataset):
         self.tokenizer = tokenizer
         self.tokenizer_kwargs = tokenizer_kwargs or {}
         self.num_samples = num_samples
+        owns_catalog = catalog is None
         if catalog is None:
             catalog = get_catalog()
         self._init_catalog(catalog)
@@ -84,7 +85,7 @@ class PytorchDataset(IterableDataset):
             self.prefetch = prefetch
 
         self._cache = catalog.cache
-        self._prefetch_cache: Optional[Cache] = None
+        self._prefetch_cache: Cache | None = None
         self._remove_prefetched = remove_prefetched
         if prefetch and not self.cache:
             tmp_dir = catalog.cache.tmp_dir
@@ -93,6 +94,10 @@ class PytorchDataset(IterableDataset):
             self._cache = self._prefetch_cache
             weakref.finalize(self, self._prefetch_cache.destroy)
 
+        # Close the catalog if we created it - we only needed it for clone params
+        if owns_catalog:
+            catalog.close()
+
     def close(self) -> None:
         if self._prefetch_cache:
             self._prefetch_cache.destroy()
@@ -104,7 +109,7 @@ class PytorchDataset(IterableDataset):
         self._ms_params = catalog.metastore.clone_params()
         self._wh_params = catalog.warehouse.clone_params()
         self._catalog_params = catalog.get_init_params()
-        self.catalog: Optional[Catalog] = None
+        self.catalog: Catalog | None = None
 
     def _get_catalog(self) -> "Catalog":
         ms_cls, ms_args, ms_kwargs = self._ms_params
@@ -121,19 +126,22 @@ class PytorchDataset(IterableDataset):
         total_workers: int,
     ) -> Generator[tuple[Any, ...], None, None]:
         catalog = self._get_catalog()
-        session = Session("PyTorch", catalog=catalog)
-        ds = read_dataset(
-            name=self.name, version=self.version, session=session
-        ).settings(cache=self.cache, prefetch=self.prefetch)
-
-        # remove file signals from dataset
-        schema = ds.signals_schema.clone_without_file_signals()
-        ds = ds.select(*schema.values.keys())
-
-        if self.num_samples > 0:
-            ds = ds.sample(self.num_samples)
-        ds = ds.chunk(total_rank, total_workers)
-        yield from ds.to_iter()
+        try:
+            session = Session("PyTorch", catalog=catalog)
+            ds = read_dataset(
+                name=self.name, version=self.version, session=session
+            ).settings(cache=self.cache, prefetch=self.prefetch)
+
+            # remove file signals from dataset
+            schema = ds.signals_schema.clone_without_file_signals()
+            ds = ds.select(*schema.values.keys())
+
+            if self.num_samples > 0:
+                ds = ds.sample(self.num_samples)
+            ds = ds.chunk(total_rank, total_workers)
+            yield from ds.to_iter()
+        finally:
+            catalog.close()
 
     def _iter_with_prefetch(self) -> Generator[tuple[Any], None, None]:
         from datachain.lib.udf import _prefetch_inputs

datachain/lib/settings.py

@@ -1,111 +1,214 @@
+from typing import Any
+
 from datachain.lib.utils import DataChainParamsError
-from datachain.utils import DEFAULT_CHUNK_ROWS
+
+DEFAULT_CACHE = False
+DEFAULT_PREFETCH = 2
+DEFAULT_BATCH_SIZE = 2_000
 
 
 class SettingsError(DataChainParamsError):
-    def __init__(self, msg):
+    def __init__(self, msg: str) -> None:
         super().__init__(f"Dataset settings error: {msg}")
 
 
 class Settings:
-    def __init__(
+    """Settings for datachain."""
+
+    _cache: bool | None
+    _prefetch: int | None
+    _parallel: bool | int | None
+    _workers: int | None
+    _namespace: str | None
+    _project: str | None
+    _min_task_size: int | None
+    _batch_size: int | None
+
+    def __init__(  # noqa: C901, PLR0912
         self,
-        cache=None,
-        parallel=None,
-        workers=None,
-        min_task_size=None,
-        prefetch=None,
-        namespace=None,
-        project=None,
-        batch_rows=None,
-    ):
-        self._cache = cache
-        self.parallel = parallel
-        self._workers = workers
-        self.min_task_size = min_task_size
-        self.prefetch = prefetch
-        self.namespace = namespace
-        self.project = project
-        self._chunk_rows = batch_rows
-
-        if not isinstance(cache, bool) and cache is not None:
-            raise SettingsError(
-                "'cache' argument must be bool"
-                f" while {cache.__class__.__name__} was given"
-            )
-
-        if not isinstance(parallel, int) and parallel is not None:
-            raise SettingsError(
-                "'parallel' argument must be int or None"
-                f" while {parallel.__class__.__name__} was given"
-            )
-
-        if (
-            not isinstance(workers, bool)
-            and not isinstance(workers, int)
-            and workers is not None
-        ):
-            raise SettingsError(
-                "'workers' argument must be int or bool"
-                f" while {workers.__class__.__name__} was given"
-            )
-
-        if min_task_size is not None and not isinstance(min_task_size, int):
-            raise SettingsError(
-                "'min_task_size' argument must be int or None"
-                f", {min_task_size.__class__.__name__} was given"
-            )
-
-        if batch_rows is not None and not isinstance(batch_rows, int):
-            raise SettingsError(
-                "'batch_rows' argument must be int or None"
-                f", {batch_rows.__class__.__name__} was given"
-            )
-
-        if batch_rows is not None and batch_rows <= 0:
-            raise SettingsError(
-                "'batch_rows' argument must be positive integer"
-                f", {batch_rows} was given"
-            )
+        cache: bool | None = None,
+        prefetch: bool | int | None = None,
+        parallel: bool | int | None = None,
+        workers: int | None = None,
+        namespace: str | None = None,
+        project: str | None = None,
+        min_task_size: int | None = None,
+        batch_size: int | None = None,
+    ) -> None:
+        if cache is None:
+            self._cache = None
+        else:
+            if not isinstance(cache, bool):
+                raise SettingsError(
+                    "'cache' argument must be bool"
+                    f" while {cache.__class__.__name__} was given"
+                )
+            self._cache = cache
+
+        if prefetch is None or prefetch is True:
+            self._prefetch = None
+        elif prefetch is False:
+            self._prefetch = 0  # disable prefetch (False == 0)
+        else:
+            if not isinstance(prefetch, int):
+                raise SettingsError(
+                    "'prefetch' argument must be int or bool"
+                    f" while {prefetch.__class__.__name__} was given"
+                )
+            if prefetch < 0:
+                raise SettingsError(
+                    "'prefetch' argument must be non-negative integer"
+                    f", {prefetch} was given"
+                )
+            self._prefetch = prefetch
+
+        if parallel is None or parallel is False:
+            self._parallel = None
+        elif parallel is True:
+            self._parallel = True
+        else:
+            if not isinstance(parallel, int):
+                raise SettingsError(
+                    "'parallel' argument must be int or bool"
+                    f" while {parallel.__class__.__name__} was given"
+                )
+            if parallel <= 0:
+                raise SettingsError(
+                    "'parallel' argument must be positive integer"
+                    f", {parallel} was given"
+                )
+            self._parallel = parallel
+
+        if workers is None:
+            self._workers = None
+        else:
+            if not isinstance(workers, int) or isinstance(workers, bool):
+                raise SettingsError(
+                    "'workers' argument must be int"
+                    f" while {workers.__class__.__name__} was given"
+                )
+            if workers <= 0:
+                raise SettingsError(
+                    f"'workers' argument must be positive integer, {workers} was given"
+                )
+            self._workers = workers
+
+        if namespace is None:
+            self._namespace = None
+        else:
+            if not isinstance(namespace, str):
+                raise SettingsError(
+                    "'namespace' argument must be str"
+                    f", {namespace.__class__.__name__} was given"
+                )
+            self._namespace = namespace
+
+        if project is None:
+            self._project = None
+        else:
+            if not isinstance(project, str):
+                raise SettingsError(
+                    "'project' argument must be str"
+                    f", {project.__class__.__name__} was given"
+                )
+            self._project = project
+
+        if min_task_size is None:
+            self._min_task_size = None
+        else:
+            if not isinstance(min_task_size, int) or isinstance(min_task_size, bool):
+                raise SettingsError(
+                    "'min_task_size' argument must be int"
+                    f", {min_task_size.__class__.__name__} was given"
+                )
+            if min_task_size <= 0:
+                raise SettingsError(
+                    "'min_task_size' argument must be positive integer"
+                    f", {min_task_size} was given"
+                )
+            self._min_task_size = min_task_size
+
+        if batch_size is None:
+            self._batch_size = None
+        else:
+            if not isinstance(batch_size, int) or isinstance(batch_size, bool):
+                raise SettingsError(
+                    "'batch_size' argument must be int"
+                    f", {batch_size.__class__.__name__} was given"
+                )
+            if batch_size <= 0:
+                raise SettingsError(
+                    "'batch_size' argument must be positive integer"
+                    f", {batch_size} was given"
+                )
+            self._batch_size = batch_size
+
+    @property
+    def cache(self) -> bool:
+        return self._cache if self._cache is not None else DEFAULT_CACHE
+
+    @property
+    def prefetch(self) -> int | None:
+        return self._prefetch if self._prefetch is not None else DEFAULT_PREFETCH
+
+    @property
+    def parallel(self) -> bool | int | None:
+        return self._parallel if self._parallel is not None else None
+
+    @property
+    def workers(self) -> int | None:
+        return self._workers if self._workers is not None else None
+
+    @property
+    def namespace(self) -> str | None:
+        return self._namespace if self._namespace is not None else None
 
     @property
-    def cache(self):
-        return self._cache if self._cache is not None else False
+    def project(self) -> str | None:
+        return self._project if self._project is not None else None
 
     @property
-    def workers(self):
-        return self._workers if self._workers is not None else False
+    def min_task_size(self) -> int | None:
+        return self._min_task_size if self._min_task_size is not None else None
 
     @property
-    def batch_rows(self):
-        return self._chunk_rows if self._chunk_rows is not None else DEFAULT_CHUNK_ROWS
+    def batch_size(self) -> int:
+        return self._batch_size if self._batch_size is not None else DEFAULT_BATCH_SIZE
 
-    def to_dict(self):
-        res = {}
+    def to_dict(self) -> dict[str, Any]:
+        res: dict[str, Any] = {}
         if self._cache is not None:
             res["cache"] = self.cache
-        if self.parallel is not None:
+        if self._prefetch is not None:
+            res["prefetch"] = self.prefetch
+        if self._parallel is not None:
             res["parallel"] = self.parallel
         if self._workers is not None:
             res["workers"] = self.workers
-        if self.min_task_size is not None:
+        if self._min_task_size is not None:
             res["min_task_size"] = self.min_task_size
-        if self.namespace is not None:
+        if self._namespace is not None:
             res["namespace"] = self.namespace
-        if self.project is not None:
+        if self._project is not None:
             res["project"] = self.project
-        if self._chunk_rows is not None:
-            res["batch_rows"] = self._chunk_rows
+        if self._batch_size is not None:
+            res["batch_size"] = self.batch_size
         return res
 
-    def add(self, settings: "Settings"):
-        self._cache = settings._cache or self._cache
-        self.parallel = settings.parallel or self.parallel
-        self._workers = settings._workers or self._workers
-        self.min_task_size = settings.min_task_size or self.min_task_size
-        self.namespace = settings.namespace or self.namespace
-        self.project = settings.project or self.project
-        if settings.prefetch is not None:
-            self.prefetch = settings.prefetch
-        if settings._chunk_rows is not None:
-            self._chunk_rows = settings._chunk_rows
+    def add(self, settings: "Settings") -> None:
+        if settings._cache is not None:
+            self._cache = settings._cache
+        if settings._prefetch is not None:
+            self._prefetch = settings._prefetch
+        if settings._parallel is not None:
+            self._parallel = settings._parallel
+        if settings._workers is not None:
+            self._workers = settings._workers
+        if settings._namespace is not None:
+            self._namespace = settings._namespace
+        if settings._project is not None:
+            self._project = settings._project
+        if settings._min_task_size is not None:
+            self._min_task_size = settings._min_task_size
+        if settings._batch_size is not None:
+            self._batch_size = settings._batch_size