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

datachain/model/ultralytics/pose.py

@@ -56,16 +56,16 @@ class YoloPose(DataModel):
         if not summary:
             return YoloPose(box=BBox(), pose=Pose3D())
         name = summary[0].get("name", "")
-        box = (
-            BBox.from_dict(summary[0]["box"], title=name)
-            if "box" in summary[0]
-            else BBox()
-        )
-        pose = (
-            Pose3D.from_dict(summary[0]["keypoints"])
-            if "keypoints" in summary[0]
-            else Pose3D()
-        )
+        if summary[0].get("box"):
+            assert isinstance(summary[0]["box"], dict)
+            box = BBox.from_dict(summary[0]["box"], title=name)
+        else:
+            box = BBox()
+        if summary[0].get("keypoints"):
+            assert isinstance(summary[0]["keypoints"], dict)
+            pose = Pose3D.from_dict(summary[0]["keypoints"])
+        else:
+            pose = Pose3D()
         return YoloPose(
             cls=summary[0]["class"],
             name=name,
@@ -102,8 +102,12 @@ class YoloPoses(DataModel):
                 cls.append(s["class"])
                 names.append(name)
                 confidence.append(s["confidence"])
-                box.append(BBox.from_dict(s.get("box", {}), title=name))
-                pose.append(Pose3D.from_dict(s.get("keypoints", {})))
+                if s.get("box"):
+                    assert isinstance(s["box"], dict)
+                    box.append(BBox.from_dict(s["box"], title=name))
+                if s.get("keypoints"):
+                    assert isinstance(s["keypoints"], dict)
+                    pose.append(Pose3D.from_dict(s["keypoints"]))
         return YoloPoses(
             cls=cls,
             name=names,

datachain/model/ultralytics/segment.py

@@ -34,16 +34,16 @@ class YoloSegment(DataModel):
         if not summary:
             return YoloSegment(box=BBox(), segment=Segment())
         name = summary[0].get("name", "")
-        box = (
-            BBox.from_dict(summary[0]["box"], title=name)
-            if "box" in summary[0]
-            else BBox()
-        )
-        segment = (
-            Segment.from_dict(summary[0]["segments"], title=name)
-            if "segments" in summary[0]
-            else Segment()
-        )
+        if summary[0].get("box"):
+            assert isinstance(summary[0]["box"], dict)
+            box = BBox.from_dict(summary[0]["box"], title=name)
+        else:
+            box = BBox()
+        if summary[0].get("segments"):
+            assert isinstance(summary[0]["segments"], dict)
+            segment = Segment.from_dict(summary[0]["segments"], title=name)
+        else:
+            segment = Segment()
         return YoloSegment(
             cls=summary[0]["class"],
             name=summary[0]["name"],
@@ -80,8 +80,12 @@ class YoloSegments(DataModel):
                 cls.append(s["class"])
                 names.append(name)
                 confidence.append(s["confidence"])
-                box.append(BBox.from_dict(s.get("box", {}), title=name))
-                segment.append(Segment.from_dict(s.get("segments", {}), title=name))
+                if s.get("box"):
+                    assert isinstance(s["box"], dict)
+                    box.append(BBox.from_dict(s["box"], title=name))
+                if s.get("segments"):
+                    assert isinstance(s["segments"], dict)
+                    segment.append(Segment.from_dict(s["segments"], title=name))
         return YoloSegments(
             cls=cls,
             name=names,

datachain/namespace.py

@@ -0,0 +1,84 @@
+import builtins
+from dataclasses import dataclass, fields
+from datetime import datetime
+from typing import Any, TypeVar
+
+from datachain.error import InvalidNamespaceNameError
+
+N = TypeVar("N", bound="Namespace")
+NAMESPACE_NAME_RESERVED_CHARS = [".", "@"]
+
+
+def parse_name(name: str) -> tuple[str, str | None]:
+    """
+    Parses namespace name into namespace and optional project name.
+    If both namespace and project are defined in name, they need to be split by dot
+    e.g dev.my-project
+    Valid inputs:
+        - dev.my-project
+        - dev
+    """
+    parts = name.split(".")
+    if len(parts) == 1:
+        return name, None
+    if len(parts) == 2:
+        return parts[0], parts[1]
+    raise InvalidNamespaceNameError(
+        f"Invalid namespace format: {name}. Expected 'namespace' or 'ns1.ns2'."
+    )
+
+
+@dataclass(frozen=True)
+class Namespace:
+    id: int
+    uuid: str
+    name: str
+    descr: str | None
+    created_at: datetime
+
+    @staticmethod
+    def validate_name(name: str) -> None:
+        """Throws exception if name is invalid, otherwise returns None"""
+        if not name:
+            raise InvalidNamespaceNameError("Namespace name cannot be empty")
+
+        for c in NAMESPACE_NAME_RESERVED_CHARS:
+            if c in name:
+                raise InvalidNamespaceNameError(
+                    f"Character {c} is reserved and not allowed in namespace name"
+                )
+
+        if name in [Namespace.default(), Namespace.system()]:
+            raise InvalidNamespaceNameError(
+                f"Namespace name {name} is reserved and cannot be used."
+            )
+
+    @staticmethod
+    def default() -> str:
+        """Name of default namespace"""
+        return "local"
+
+    @staticmethod
+    def system() -> str:
+        """Name of the system namespace"""
+        return "system"
+
+    @property
+    def is_system(self):
+        return self.name == Namespace.system()
+
+    @classmethod
+    def parse(
+        cls: builtins.type[N],
+        id: int,
+        uuid: str,
+        name: str,
+        descr: str | None,
+        created_at: datetime,
+    ) -> "Namespace":
+        return cls(id, uuid, name, descr, created_at)
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Namespace":
+        kwargs = {f.name: d[f.name] for f in fields(cls) if f.name in d}
+        return cls(**kwargs)

datachain/node.py

@@ -1,6 +1,6 @@
 import os
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 import attrs
 
@@ -53,11 +53,11 @@ class Node:
     sys__rand: int = 0
     path: str = ""
     etag: str = ""
-    version: Optional[str] = None
+    version: str | None = None
     is_latest: bool = True
-    last_modified: Optional[datetime] = None
+    last_modified: datetime | None = None
     size: int = 0
-    location: Optional[str] = None
+    location: str | None = None
     source: StorageURI = StorageURI("")  # noqa: RUF009
     dir_type: int = DirType.FILE
 
@@ -90,7 +90,7 @@ class Node:
             return self.path + "/"
         return self.path
 
-    def to_file(self, source: Optional[StorageURI] = None) -> File:
+    def to_file(self, source: StorageURI | None = None) -> File:
         if source is None:
             source = self.source
         return File(
@@ -189,7 +189,7 @@ class NodeWithPath:
 TIME_FMT = "%Y-%m-%d %H:%M"
 
 
-def long_line_str(name: str, timestamp: Optional[datetime]) -> str:
+def long_line_str(name: str, timestamp: datetime | None) -> str:
     if timestamp is None:
         time = "-"
     else:

datachain/nodes_thread_pool.py

@@ -1,4 +1,3 @@
-import concurrent
 import concurrent.futures
 import threading
 from abc import ABC, abstractmethod

datachain/plugins.py

@@ -0,0 +1,24 @@
+"""Plugin loader for DataChain callables.
+
+Discovers and invokes entry points in the group "datachain.callables" once
+per process. This enables external packages (e.g., Studio) to register
+their callables with the serializer registry without explicit imports.
+"""
+
+from importlib import metadata as importlib_metadata
+
+_plugins_loaded = False
+
+
+def ensure_plugins_loaded() -> None:
+    global _plugins_loaded  # noqa: PLW0603
+    if _plugins_loaded:
+        return
+
+    # Compatible across importlib.metadata versions
+    eps_obj = importlib_metadata.entry_points()
+    for ep in eps_obj.select(group="datachain.callables"):
+        func = ep.load()
+        func()
+
+    _plugins_loaded = True

datachain/project.py

@@ -0,0 +1,78 @@
+import builtins
+from dataclasses import dataclass, fields
+from datetime import datetime
+from typing import Any, TypeVar
+
+from datachain.error import InvalidProjectNameError
+from datachain.namespace import Namespace
+
+P = TypeVar("P", bound="Project")
+PROJECT_NAME_RESERVED_CHARS = [".", "@"]
+
+
+@dataclass(frozen=True)
+class Project:
+    id: int
+    uuid: str
+    name: str
+    descr: str | None
+    created_at: datetime
+    namespace: Namespace
+
+    @staticmethod
+    def validate_name(name: str) -> None:
+        """Throws exception if name is invalid, otherwise returns None"""
+        if not name:
+            raise InvalidProjectNameError("Project name cannot be empty")
+
+        for c in PROJECT_NAME_RESERVED_CHARS:
+            if c in name:
+                raise InvalidProjectNameError(
+                    f"Character {c} is reserved and not allowed in project name."
+                )
+
+        if name in [Project.default(), Project.listing()]:
+            raise InvalidProjectNameError(
+                f"Project name {name} is reserved and cannot be used."
+            )
+
+    @staticmethod
+    def default() -> str:
+        """Name of default project"""
+        return "local"
+
+    @staticmethod
+    def listing() -> str:
+        """Name of listing project where all listing datasets will be saved"""
+        return "listing"
+
+    @classmethod
+    def parse(
+        cls: builtins.type[P],
+        namespace_id: int,
+        namespace_uuid: str,
+        namespace_name: str,
+        namespace_descr: str | None,
+        namespace_created_at: datetime,
+        project_id: int,
+        uuid: str,
+        name: str,
+        descr: str | None,
+        created_at: datetime,
+        project_namespace_id: int,
+    ) -> "Project":
+        namespace = Namespace.parse(
+            namespace_id,
+            namespace_uuid,
+            namespace_name,
+            namespace_descr,
+            namespace_created_at,
+        )
+
+        return cls(project_id, uuid, name, descr, created_at, namespace)
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Project":
+        namespace = Namespace.from_dict(d.pop("namespace"))
+        kwargs = {f.name: d[f.name] for f in fields(cls) if f.name in d}
+        return cls(**kwargs, namespace=namespace)

datachain/query/batch.py

@@ -1,24 +1,14 @@
 import contextlib
 import math
 from abc import ABC, abstractmethod
-from collections.abc import Generator, Sequence
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, Optional, Union
-
-from datachain.data_storage.schema import PARTITION_COLUMN_ID
-from datachain.data_storage.warehouse import SELECT_BATCH_SIZE
-from datachain.query.utils import get_query_column, get_query_id_column
-
-if TYPE_CHECKING:
-    from sqlalchemy import Select
+from collections.abc import Callable, Generator, Sequence
 
+import sqlalchemy as sa
 
-@dataclass
-class RowsOutputBatch:
-    rows: Sequence[Sequence]
-
+from datachain.data_storage.schema import PARTITION_COLUMN_ID
 
-RowsOutput = Union[Sequence, RowsOutputBatch]
+RowsOutputBatch = Sequence[Sequence]
+RowsOutput = Sequence | RowsOutputBatch
 
 
 class BatchingStrategy(ABC):
@@ -30,8 +20,8 @@ class BatchingStrategy(ABC):
     def __call__(
         self,
         execute: Callable,
-        query: "Select",
-        ids_only: bool = False,
+        query: sa.Select,
+        id_col: sa.ColumnElement | None = None,
     ) -> Generator[RowsOutput, None, None]:
         """Apply the provided parameters to the UDF."""
 
@@ -47,12 +37,16 @@ class NoBatching(BatchingStrategy):
     def __call__(
         self,
         execute: Callable,
-        query: "Select",
-        ids_only: bool = False,
+        query: sa.Select,
+        id_col: sa.ColumnElement | None = None,
     ) -> Generator[Sequence, None, None]:
-        if ids_only:
-            query = query.with_only_columns(get_query_id_column(query))
-        return execute(query)
+        ids_only = False
+        if id_col is not None:
+            query = query.with_only_columns(id_col)
+            ids_only = True
+
+        rows = execute(query)
+        yield from (r[0] for r in rows) if ids_only else rows
 
 
 class Batch(BatchingStrategy):
@@ -69,27 +63,31 @@ class Batch(BatchingStrategy):
     def __call__(
         self,
         execute: Callable,
-        query: "Select",
-        ids_only: bool = False,
-    ) -> Generator[RowsOutputBatch, None, None]:
-        if ids_only:
-            query = query.with_only_columns(get_query_id_column(query))
+        query: sa.Select,
+        id_col: sa.ColumnElement | None = None,
+    ) -> Generator[RowsOutput, None, None]:
+        from datachain.data_storage.warehouse import SELECT_BATCH_SIZE
+
+        ids_only = False
+        if id_col is not None:
+            query = query.with_only_columns(id_col)
+            ids_only = True
 
         # choose page size that is a multiple of the batch size
         page_size = math.ceil(SELECT_BATCH_SIZE / self.count) * self.count
 
         # select rows in batches
-        results: list[Sequence] = []
+        results = []
 
         with contextlib.closing(execute(query, page_size=page_size)) as rows:
             for row in rows:
                 results.append(row)
                 if len(results) >= self.count:
                     batch, results = results[: self.count], results[self.count :]
-                    yield RowsOutputBatch(batch)
+                    yield [r[0] for r in batch] if ids_only else batch
 
             if len(results) > 0:
-                yield RowsOutputBatch(results)
+                yield [r[0] for r in results] if ids_only else results
 
 
 class Partition(BatchingStrategy):
@@ -104,18 +102,19 @@ class Partition(BatchingStrategy):
     def __call__(
         self,
         execute: Callable,
-        query: "Select",
-        ids_only: bool = False,
-    ) -> Generator[RowsOutputBatch, None, None]:
-        id_col = get_query_id_column(query)
-        if (partition_col := get_query_column(query, PARTITION_COLUMN_ID)) is None:
+        query: sa.Select,
+        id_col: sa.ColumnElement | None = None,
+    ) -> Generator[RowsOutput, None, None]:
+        if (partition_col := query.selected_columns.get(PARTITION_COLUMN_ID)) is None:
             raise RuntimeError("partition column not found in query")
 
-        if ids_only:
+        ids_only = False
+        if id_col is not None:
             query = query.with_only_columns(id_col, partition_col)
+            ids_only = True
 
-        current_partition: Optional[int] = None
-        batch: list[Sequence] = []
+        current_partition: int | None = None
+        batch: list = []
 
         query_fields = [str(c.name) for c in query.selected_columns]
         id_column_idx = query_fields.index("sys__id")
@@ -132,9 +131,9 @@ class Partition(BatchingStrategy):
                 if current_partition != partition:
                     current_partition = partition
                     if len(batch) > 0:
-                        yield RowsOutputBatch(batch)
+                        yield batch
                         batch = []
-                batch.append([row[id_column_idx]] if ids_only else row)
+                batch.append(row[id_column_idx] if ids_only else row)
 
             if len(batch) > 0:
-                yield RowsOutputBatch(batch)
+                yield batch