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

datachain/lib/webdataset_laion.py

@@ -1,6 +1,5 @@
 import warnings
 from collections.abc import Iterator
-from typing import Optional
 
 import numpy as np
 from pydantic import BaseModel, Field
@@ -23,18 +22,18 @@ warnings.filterwarnings(
 
 class Laion(WDSReadableSubclass):
     uid: str = Field(default="")
-    face_bboxes: Optional[list[list[float]]] = Field(default=None)
-    caption: Optional[str] = Field(default=None)
-    url: Optional[str] = Field(default=None)
-    key: Optional[str] = Field(default=None)
-    status: Optional[str] = Field(default=None)
-    error_message: Optional[str] = Field(default=None)
-    width: Optional[int] = Field(default=None)
-    height: Optional[int] = Field(default=None)
-    original_width: Optional[int] = Field(default=None)
-    original_height: Optional[int] = Field(default=None)
-    exif: Optional[str] = Field(default=None)
-    sha256: Optional[str] = Field(default=None)
+    face_bboxes: list[list[float]] | None = Field(default=None)
+    caption: str | None = Field(default=None)
+    url: str | None = Field(default=None)
+    key: str | None = Field(default=None)
+    status: str | None = Field(default=None)
+    error_message: str | None = Field(default=None)
+    width: int | None = Field(default=None)
+    height: int | None = Field(default=None)
+    original_width: int | None = Field(default=None)
+    original_height: int | None = Field(default=None)
+    exif: str | None = Field(default=None)
+    sha256: str | None = Field(default=None)
 
     @staticmethod
     def _reader(builder, item):
@@ -42,13 +41,13 @@ class Laion(WDSReadableSubclass):
 
 
 class WDSLaion(WDSBasic):
-    txt: Optional[str] = Field(default=None)
-    json: Laion  # type: ignore[assignment]
+    txt: str | None = Field(default=None)
+    json: Laion = Field(default_factory=Laion)  # type: ignore[assignment]
 
 
 class LaionMeta(BaseModel):
     file: File
-    index: Optional[int] = Field(default=None)
+    index: int | None = Field(default=None)
     b32_img: list[float] = Field(default=[])
     b32_txt: list[float] = Field(default=[])
     l14_img: list[float] = Field(default=[])

datachain/listing.py

@@ -2,7 +2,7 @@ import glob
 import os
 from collections.abc import Iterable, Iterator
 from functools import cached_property
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 from sqlalchemy import Column
 from sqlalchemy.sql import func
@@ -25,8 +25,8 @@ class Listing:
         metastore: "AbstractMetastore",
         warehouse: "AbstractWarehouse",
         client: "Client",
-        dataset_name: Optional["str"] = None,
-        dataset_version: Optional[str] = None,
+        dataset_name: str | None = None,
+        dataset_version: str | None = None,
         column: str = "file",
     ):
         self.metastore = metastore
@@ -35,6 +35,7 @@ class Listing:
         self.dataset_name = dataset_name  # dataset representing bucket listing
         self.dataset_version = dataset_version  # dataset representing bucket listing
         self.column = column
+        self._closed = False
 
     def clone(self) -> "Listing":
         return self.__class__(
@@ -53,7 +54,13 @@ class Listing:
         self.close()
 
     def close(self) -> None:
-        self.warehouse.close()
+        if self._closed:
+            return
+        self._closed = True
+        try:
+            self.warehouse.close_on_exit()
+        finally:
+            self.metastore.close_on_exit()
 
     @property
     def uri(self):
@@ -102,7 +109,7 @@ class Listing:
     def collect_nodes_to_instantiate(
         self,
         sources: Iterable["DataSource"],
-        copy_to_filename: Optional[str],
+        copy_to_filename: str | None,
         recursive=False,
         copy_dir_contents=False,
         from_dataset=False,

datachain/model/bbox.py

@@ -198,7 +198,9 @@ class BBox(DataModel):
     def pose_inside(self, pose: Union["Pose", "Pose3D"]) -> bool:
         """Return True if the pose is inside the bounding box."""
         return all(
-            self.point_inside(x, y) for x, y in zip(pose.x, pose.y) if x > 0 or y > 0
+            self.point_inside(x, y)
+            for x, y in zip(pose.x, pose.y, strict=False)
+            if x > 0 or y > 0
         )
 
     @staticmethod

datachain/namespace.py

@@ -1,7 +1,7 @@
 import builtins
 from dataclasses import dataclass, fields
 from datetime import datetime
-from typing import Any, Optional, TypeVar
+from typing import Any, TypeVar
 
 from datachain.error import InvalidNamespaceNameError
 
@@ -9,12 +9,31 @@ 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: Optional[str]
+    descr: str | None
     created_at: datetime
 
     @staticmethod
@@ -54,7 +73,7 @@ class Namespace:
         id: int,
         uuid: str,
         name: str,
-        descr: Optional[str],
+        descr: str | None,
         created_at: datetime,
     ) -> "Namespace":
         return cls(id, uuid, name, descr, created_at)

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

@@ -1,7 +1,7 @@
 import builtins
 from dataclasses import dataclass, fields
 from datetime import datetime
-from typing import Any, Optional, TypeVar
+from typing import Any, TypeVar
 
 from datachain.error import InvalidProjectNameError
 from datachain.namespace import Namespace
@@ -15,7 +15,7 @@ class Project:
     id: int
     uuid: str
     name: str
-    descr: Optional[str]
+    descr: str | None
     created_at: datetime
     namespace: Namespace
 
@@ -52,12 +52,12 @@ class Project:
         namespace_id: int,
         namespace_uuid: str,
         namespace_name: str,
-        namespace_descr: Optional[str],
+        namespace_descr: str | None,
         namespace_created_at: datetime,
         project_id: int,
         uuid: str,
         name: str,
-        descr: Optional[str],
+        descr: str | None,
         created_at: datetime,
         project_namespace_id: int,
     ) -> "Project":

datachain/query/batch.py

@@ -1,16 +1,14 @@
 import contextlib
 import math
 from abc import ABC, abstractmethod
-from collections.abc import Generator, Sequence
-from typing import Callable, Optional, Union
+from collections.abc import Callable, Generator, Sequence
 
 import sqlalchemy as sa
 
 from datachain.data_storage.schema import PARTITION_COLUMN_ID
-from datachain.query.utils import get_query_column
 
 RowsOutputBatch = Sequence[Sequence]
-RowsOutput = Union[Sequence, RowsOutputBatch]
+RowsOutput = Sequence | RowsOutputBatch
 
 
 class BatchingStrategy(ABC):
@@ -23,7 +21,7 @@ class BatchingStrategy(ABC):
         self,
         execute: Callable,
         query: sa.Select,
-        id_col: Optional[sa.ColumnElement] = None,
+        id_col: sa.ColumnElement | None = None,
     ) -> Generator[RowsOutput, None, None]:
         """Apply the provided parameters to the UDF."""
 
@@ -40,7 +38,7 @@ class NoBatching(BatchingStrategy):
         self,
         execute: Callable,
         query: sa.Select,
-        id_col: Optional[sa.ColumnElement] = None,
+        id_col: sa.ColumnElement | None = None,
     ) -> Generator[Sequence, None, None]:
         ids_only = False
         if id_col is not None:
@@ -66,7 +64,7 @@ class Batch(BatchingStrategy):
         self,
         execute: Callable,
         query: sa.Select,
-        id_col: Optional[sa.ColumnElement] = None,
+        id_col: sa.ColumnElement | None = None,
     ) -> Generator[RowsOutput, None, None]:
         from datachain.data_storage.warehouse import SELECT_BATCH_SIZE
 
@@ -81,8 +79,8 @@ class Batch(BatchingStrategy):
         # select rows in batches
         results = []
 
-        with contextlib.closing(execute(query, page_size=page_size)) as batch_rows:
-            for row in batch_rows:
+        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 :]
@@ -105,9 +103,9 @@ class Partition(BatchingStrategy):
         self,
         execute: Callable,
         query: sa.Select,
-        id_col: Optional[sa.ColumnElement] = None,
+        id_col: sa.ColumnElement | None = None,
     ) -> Generator[RowsOutput, None, None]:
-        if (partition_col := get_query_column(query, PARTITION_COLUMN_ID)) is None:
+        if (partition_col := query.selected_columns.get(PARTITION_COLUMN_ID)) is None:
             raise RuntimeError("partition column not found in query")
 
         ids_only = False
@@ -115,7 +113,7 @@ class Partition(BatchingStrategy):
             query = query.with_only_columns(id_col, partition_col)
             ids_only = True
 
-        current_partition: Optional[int] = None
+        current_partition: int | None = None
         batch: list = []
 
         query_fields = [str(c.name) for c in query.selected_columns]