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

datachain/lib/utils.py

@@ -1,6 +1,9 @@
+import inspect
 import re
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
+from pathlib import PurePosixPath
+from urllib.parse import urlparse
 
 
 class AbstractUDF(ABC):
@@ -18,13 +21,11 @@ class AbstractUDF(ABC):
 
 
 class DataChainError(Exception):
-    def __init__(self, message):
-        super().__init__(message)
+    pass
 
 
 class DataChainParamsError(DataChainError):
-    def __init__(self, message):
-        super().__init__(message)
+    pass
 
 
 class DataChainColumnError(DataChainParamsError):
@@ -32,6 +33,25 @@ class DataChainColumnError(DataChainParamsError):
         super().__init__(f"Error for column {col_name}: {msg}")
 
 
+def callable_name(obj: object) -> str:
+    """Return a friendly name for a callable or UDF-like instance."""
+    # UDF classes in DataChain inherit from AbstractUDF; prefer class name
+    if isinstance(obj, AbstractUDF):
+        return obj.__class__.__name__
+
+    # Plain functions and bound/unbound methods
+    if inspect.ismethod(obj) or inspect.isfunction(obj):
+        # __name__ exists for functions/methods; includes "<lambda>" for lambdas
+        return obj.__name__  # type: ignore[attr-defined]
+
+    # Generic callable object
+    if callable(obj):
+        return obj.__class__.__name__
+
+    # Fallback for non-callables
+    return str(obj)
+
+
 def normalize_col_names(col_names: Sequence[str]) -> dict[str, str]:
     """Returns normalized_name -> original_name dict."""
     gen_col_counter = 0
@@ -59,3 +79,97 @@ def normalize_col_names(col_names: Sequence[str]) -> dict[str, str]:
         new_col_names[generated_column] = org_column
 
     return new_col_names
+
+
+def rebase_path(
+    src_path: str,
+    old_base: str,
+    new_base: str,
+    suffix: str = "",
+    extension: str = "",
+) -> str:
+    """
+    Rebase a file path from one base directory to another.
+
+    Args:
+        src_path: Source file path (can include URI scheme like s3://)
+        old_base: Base directory to remove from src_path
+        new_base: New base directory to prepend
+        suffix: Optional suffix to add before file extension
+        extension: Optional new file extension (without dot)
+
+    Returns:
+        str: Rebased path with new base directory
+
+    Raises:
+        ValueError: If old_base is not found in src_path
+    """
+    # Parse URIs to handle schemes properly
+    src_parsed = urlparse(src_path)
+    old_base_parsed = urlparse(old_base)
+    new_base_parsed = urlparse(new_base)
+
+    # Get the path component (without scheme)
+    if src_parsed.scheme:
+        src_path_only = src_parsed.netloc + src_parsed.path
+    else:
+        src_path_only = src_path
+
+    if old_base_parsed.scheme:
+        old_base_only = old_base_parsed.netloc + old_base_parsed.path
+    else:
+        old_base_only = old_base
+
+    # Normalize paths
+    src_path_norm = PurePosixPath(src_path_only).as_posix()
+    old_base_norm = PurePosixPath(old_base_only).as_posix()
+
+    # Find where old_base appears in src_path
+    if old_base_norm in src_path_norm:
+        # Find the index where old_base appears
+        idx = src_path_norm.find(old_base_norm)
+        if idx == -1:
+            raise ValueError(f"old_base '{old_base}' not found in src_path")
+
+        # Extract the relative path after old_base
+        relative_start = idx + len(old_base_norm)
+        # Skip leading slash if present
+        if relative_start < len(src_path_norm) and src_path_norm[relative_start] == "/":
+            relative_start += 1
+        relative_path = src_path_norm[relative_start:]
+    else:
+        raise ValueError(f"old_base '{old_base}' not found in src_path")
+
+    # Parse the filename
+    path_obj = PurePosixPath(relative_path)
+    stem = path_obj.stem
+    current_ext = path_obj.suffix
+
+    # Apply suffix and extension changes
+    new_stem = stem + suffix if suffix else stem
+    if extension:
+        new_ext = f".{extension}"
+    elif current_ext:
+        new_ext = current_ext
+    else:
+        new_ext = ""
+
+    # Build new filename
+    new_name = new_stem + new_ext
+
+    # Reconstruct path with new base
+    parent = str(path_obj.parent)
+    if parent == ".":
+        new_relative_path = new_name
+    else:
+        new_relative_path = str(PurePosixPath(parent) / new_name)
+
+    # Handle new_base URI scheme
+    if new_base_parsed.scheme:
+        # Has schema like s3://
+        base_path = new_base_parsed.netloc + new_base_parsed.path
+        base_path = PurePosixPath(base_path).as_posix()
+        full_path = str(PurePosixPath(base_path) / new_relative_path)
+        return f"{new_base_parsed.scheme}://{full_path}"
+    # Regular path
+    return str(PurePosixPath(new_base) / new_relative_path)

datachain/lib/video.py

@@ -1,7 +1,6 @@
 import posixpath
 import shutil
 import tempfile
-from typing import Optional, Union
 
 from numpy import ndarray
 
@@ -18,7 +17,7 @@ except ImportError as exc:
     ) from exc
 
 
-def video_info(file: Union[File, VideoFile]) -> Video:
+def video_info(file: File | VideoFile) -> Video:
     """
     Returns video file information.
 
@@ -34,21 +33,27 @@ def video_info(file: Union[File, VideoFile]) -> Video:
         file.ensure_cached()
         file_path = file.get_local_path()
         if not file_path:
-            raise FileError(file, "unable to download video file")
+            raise FileError("unable to download video file", file.source, file.path)
 
     try:
         probe = ffmpeg.probe(file_path)
     except Exception as exc:
-        raise FileError(file, "unable to extract metadata from video file") from exc
+        raise FileError(
+            "unable to extract metadata from video file", file.source, file.path
+        ) from exc
 
     all_streams = probe.get("streams")
     video_format = probe.get("format")
     if not all_streams or not video_format:
-        raise FileError(file, "unable to extract metadata from video file")
+        raise FileError(
+            "unable to extract metadata from video file", file.source, file.path
+        )
 
     video_streams = [s for s in all_streams if s["codec_type"] == "video"]
     if len(video_streams) == 0:
-        raise FileError(file, "unable to extract metadata from video file")
+        raise FileError(
+            "unable to extract metadata from video file", file.source, file.path
+        )
 
     video_stream = video_streams[0]
 
@@ -102,7 +107,7 @@ def video_frame_np(video: VideoFile, frame: int) -> ndarray:
 def validate_frame_range(
     video: VideoFile,
     start: int = 0,
-    end: Optional[int] = None,
+    end: int | None = None,
     step: int = 1,
 ) -> tuple[int, int, int]:
     """
@@ -180,7 +185,7 @@ def save_video_fragment(
     start: float,
     end: float,
     output: str,
-    format: Optional[str] = None,
+    format: str | None = None,
 ) -> VideoFile:
     """
     Saves video interval as a new video file. If output is a remote path,
@@ -199,7 +204,10 @@ def save_video_fragment(
         VideoFile: Video fragment model.
     """
     if start < 0 or end < 0 or start >= end:
-        raise ValueError(f"Invalid time range: ({start:.3f}, {end:.3f})")
+        raise ValueError(
+            f"Can't save video fragment for '{video.path}', "
+            f"invalid time range: ({start:.3f}, {end:.3f})"
+        )
 
     if format is None:
         format = video.get_file_ext()

datachain/lib/webdataset.py

@@ -1,20 +1,13 @@
-import json
 import tarfile
+import types
 import warnings
-from collections.abc import Iterator, Sequence
+from collections.abc import Callable, Iterator, Sequence
 from pathlib import Path
-from typing import (
-    Any,
-    Callable,
-    ClassVar,
-    Optional,
-    Union,
-    get_args,
-    get_origin,
-)
+from typing import Any, ClassVar, Union, get_args, get_origin
 
 from pydantic import Field
 
+from datachain import json
 from datachain.lib.data_model import DataModel
 from datachain.lib.file import File
 from datachain.lib.tar import build_tar_member
@@ -34,29 +27,29 @@ warnings.filterwarnings(
 
 
 class WDSError(DataChainError):
-    def __init__(self, tar_stream, message: str):
-        super().__init__(f"WebDataset error '{tar_stream.get_full_name()}': {message}")
+    def __init__(self, tar_name: str, message: str):
+        super().__init__(f"WebDataset error '{tar_name}': {message}")
 
 
 class CoreFileDuplicationError(WDSError):
-    def __init__(self, tar_stream, file1: str, file2: str):
+    def __init__(self, tar_name: str, file1: str, file2: str):
         super().__init__(
-            tar_stream, f"duplication of files with core extensions: {file1}, {file2}"
+            tar_name, f"duplication of files with core extensions: {file1}, {file2}"
         )
 
 
 class CoreFileNotFoundError(WDSError):
-    def __init__(self, tar_stream, extensions, stem):
+    def __init__(self, tar_name: str, extensions: Sequence[str], stem: str):
         super().__init__(
-            tar_stream,
+            tar_name,
             f"no files with the extensions '{','.join(extensions)}'"
             f" were found for file stem {stem}",
         )
 
 
 class UnknownFileExtensionError(WDSError):
-    def __init__(self, tar_stream, name, ext):
-        super().__init__(tar_stream, f"unknown extension '{ext}' for file '{name}'")
+    def __init__(self, tar_name, name: str, ext: str):
+        super().__init__(tar_name, f"unknown extension '{ext}' for file '{name}'")
 
 
 class WDSBasic(DataModel):
@@ -64,28 +57,28 @@ class WDSBasic(DataModel):
 
 
 class WDSAllFile(WDSBasic):
-    txt: Optional[str] = Field(default=None)
-    text: Optional[str] = Field(default=None)
-    cap: Optional[str] = Field(default=None)
-    transcript: Optional[str] = Field(default=None)
-    cls: Optional[int] = Field(default=None)
-    cls2: Optional[int] = Field(default=None)
-    index: Optional[int] = Field(default=None)
-    inx: Optional[int] = Field(default=None)
-    id: Optional[int] = Field(default=None)
-    json: Optional[dict] = Field(default=None)  # type: ignore[assignment]
-    jsn: Optional[dict] = Field(default=None)
-
-    pyd: Optional[bytes] = Field(default=None)
-    pickle: Optional[bytes] = Field(default=None)
-    pth: Optional[bytes] = Field(default=None)
-    ten: Optional[bytes] = Field(default=None)
-    tb: Optional[bytes] = Field(default=None)
-    mp: Optional[bytes] = Field(default=None)
-    msg: Optional[bytes] = Field(default=None)
-    npy: Optional[bytes] = Field(default=None)
-    npz: Optional[bytes] = Field(default=None)
-    cbor: Optional[bytes] = Field(default=None)
+    txt: str | None = Field(default=None)
+    text: str | None = Field(default=None)
+    cap: str | None = Field(default=None)
+    transcript: str | None = Field(default=None)
+    cls: int | None = Field(default=None)
+    cls2: int | None = Field(default=None)
+    index: int | None = Field(default=None)
+    inx: int | None = Field(default=None)
+    id: int | None = Field(default=None)
+    json: dict | None = Field(default=None)  # type: ignore[assignment]
+    jsn: dict | None = Field(default=None)
+
+    pyd: bytes | None = Field(default=None)
+    pickle: bytes | None = Field(default=None)
+    pth: bytes | None = Field(default=None)
+    ten: bytes | None = Field(default=None)
+    tb: bytes | None = Field(default=None)
+    mp: bytes | None = Field(default=None)
+    msg: bytes | None = Field(default=None)
+    npy: bytes | None = Field(default=None)
+    npz: bytes | None = Field(default=None)
+    cbor: bytes | None = Field(default=None)
 
 
 class WDSReadableSubclass(DataModel):
@@ -113,10 +106,10 @@ class Builder:
     def __init__(
         self,
         tar_stream: File,
-        core_extensions: list[str],
+        core_extensions: Sequence[str],
         wds_class: type[WDSBasic],
-        tar,
-        encoding="utf-8",
+        tar: tarfile.TarFile,
+        encoding: str = "utf-8",
     ):
         self._core_extensions = core_extensions
         self._tar_stream = tar_stream
@@ -145,18 +138,20 @@ class Builder:
         if ext in self._core_extensions:
             if self.state.core_file is not None:
                 raise CoreFileDuplicationError(
-                    self._tar_stream, file.name, self.state.core_file.name
+                    self._tar_stream.name, file.name, self.state.core_file.name
                 )
             self.state.core_file = file
         elif ext in self.state.data:
             raise WDSError(
-                self._tar_stream,
+                self._tar_stream.name,
                 f"file with extension '.{ext}' already exists in the archive",
             )
         else:
             type_ = self._get_type(ext)
             if type_ is None:
-                raise UnknownFileExtensionError(self._tar_stream, fstream.name, ext)
+                raise UnknownFileExtensionError(
+                    self._tar_stream.name, fstream.name, ext
+                )
 
             if issubclass(type_, WDSReadableSubclass):
                 reader = type_._reader
@@ -165,7 +160,7 @@ class Builder:
 
             if reader is None:
                 raise WDSError(
-                    self._tar_stream,
+                    self._tar_stream.name,
                     f"unable to find a reader for type {type_}, extension .{ext}",
                 )
             self.state.data[ext] = reader(self, file)
@@ -173,7 +168,7 @@ class Builder:
     def produce(self):
         if self.state.core_file is None:
             raise CoreFileNotFoundError(
-                self._tar_stream, self._core_extensions, self.state.stem
+                self._tar_stream.name, self._core_extensions, self.state.stem
             )
 
         file = build_tar_member(self._tar_stream, self.state.core_file)
@@ -187,14 +182,22 @@ class Builder:
             return
 
         anno = field.annotation
-        if get_origin(anno) == Union:
-            args = get_args(anno)
-            anno = args[0]
+        anno_origin = get_origin(anno)
+        if anno_origin in (Union, types.UnionType):
+            anno_args = get_args(anno)
+            if len(anno_args) == 2 and type(None) in anno_args:
+                return anno_args[0] if anno_args[1] is type(None) else anno_args[1]
 
         return anno
 
 
-def get_tar_groups(stream, tar, core_extensions, spec, encoding="utf-8"):
+def get_tar_groups(
+    stream: File,
+    tar: tarfile.TarFile,
+    core_extensions: Sequence[str],
+    spec: type[WDSBasic],
+    encoding: str = "utf-8",
+) -> Iterator[WDSBasic]:
     builder = Builder(stream, core_extensions, spec, tar, encoding)
 
     for item in sorted(tar.getmembers(), key=lambda m: Path(m.name).stem):
@@ -210,9 +213,11 @@ def get_tar_groups(stream, tar, core_extensions, spec, encoding="utf-8"):
 
 
 def process_webdataset(
-    core_extensions: Sequence[str] = ("jpg", "png"), spec=WDSAllFile, encoding="utf-8"
-) -> Callable:
-    def wds_func(file: File) -> Iterator[spec]:
+    core_extensions: Sequence[str] = ("jpg", "png"),
+    spec: type[WDSBasic] = WDSAllFile,
+    encoding: str = "utf-8",
+) -> Callable[[File], Iterator]:
+    def wds_func(file: File) -> Iterator[spec]:  # type: ignore[valid-type]
         with file.open() as fd:
             with tarfile.open(fileobj=fd) as tar:
                 yield from get_tar_groups(file, tar, core_extensions, spec, encoding)

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,16 +25,17 @@ class Listing:
         metastore: "AbstractMetastore",
         warehouse: "AbstractWarehouse",
         client: "Client",
-        dataset_name: Optional["str"] = None,
-        dataset_version: Optional[int] = None,
-        object_name: str = "file",
+        dataset_name: str | None = None,
+        dataset_version: str | None = None,
+        column: str = "file",
     ):
         self.metastore = metastore
         self.warehouse = warehouse
         self.client = client
         self.dataset_name = dataset_name  # dataset representing bucket listing
         self.dataset_version = dataset_version  # dataset representing bucket listing
-        self.object_name = object_name
+        self.column = column
+        self._closed = False
 
     def clone(self) -> "Listing":
         return self.__class__(
@@ -43,7 +44,7 @@ class Listing:
             self.client,
             self.dataset_name,
             self.dataset_version,
-            self.object_name,
+            self.column,
         )
 
     def __enter__(self) -> "Listing":
@@ -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):
@@ -66,7 +73,12 @@ class Listing:
     @cached_property
     def dataset(self) -> "DatasetRecord":
         assert self.dataset_name
-        return self.metastore.get_dataset(self.dataset_name)
+        project = self.metastore.listing_project
+        return self.metastore.get_dataset(
+            self.dataset_name,
+            namespace_name=project.namespace.name,
+            project_name=project.name,
+        )
 
     @cached_property
     def dataset_rows(self):
@@ -74,7 +86,7 @@ class Listing:
         return self.warehouse.dataset_rows(
             dataset,
             self.dataset_version or dataset.latest_version,
-            object_name=self.object_name,
+            column=self.column,
         )
 
     def expand_path(self, path, use_glob=True) -> list[Node]:
@@ -97,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/model/ultralytics/bbox.py

@@ -31,11 +31,11 @@ class YoloBBox(DataModel):
         if not summary:
             return YoloBBox(box=BBox())
         name = summary[0].get("name", "")
-        box = (
-            BBox.from_dict(summary[0]["box"], title=name)
-            if "box" in summary[0]
-            else BBox()
-        )
+        if summary[0].get("box"):
+            assert isinstance(summary[0]["box"], dict)
+            box = BBox.from_dict(summary[0]["box"], title=name)
+        else:
+            box = BBox()
         return YoloBBox(
             cls=summary[0]["class"],
             name=name,
@@ -69,7 +69,9 @@ class YoloBBoxes(DataModel):
                 cls.append(s["class"])
                 names.append(name)
                 confidence.append(s["confidence"])
-                box.append(BBox.from_dict(s.get("box", {}), title=name))
+                if s.get("box"):
+                    assert isinstance(s["box"], dict)
+                    box.append(BBox.from_dict(s["box"], title=name))
         return YoloBBoxes(
             cls=cls,
             name=names,
@@ -100,11 +102,11 @@ class YoloOBBox(DataModel):
         if not summary:
             return YoloOBBox(box=OBBox())
         name = summary[0].get("name", "")
-        box = (
-            OBBox.from_dict(summary[0]["box"], title=name)
-            if "box" in summary[0]
-            else OBBox()
-        )
+        if summary[0].get("box"):
+            assert isinstance(summary[0]["box"], dict)
+            box = OBBox.from_dict(summary[0]["box"], title=name)
+        else:
+            box = OBBox()
         return YoloOBBox(
             cls=summary[0]["class"],
             name=name,
@@ -138,7 +140,9 @@ class YoloOBBoxes(DataModel):
                 cls.append(s["class"])
                 names.append(name)
                 confidence.append(s["confidence"])
-                box.append(OBBox.from_dict(s.get("box", {}), title=name))
+                if s.get("box"):
+                    assert isinstance(s["box"], dict)
+                    box.append(OBBox.from_dict(s["box"], title=name))
         return YoloOBBoxes(
             cls=cls,
             name=names,