datachain 0.34.6__py3-none-any.whl → 0.35.0__py3-none-any.whl

datachain/lib/dc/pandas.py

@@ -1,7 +1,4 @@
-from typing import (
-    TYPE_CHECKING,
-    Optional,
-)
+from typing import TYPE_CHECKING
 
 from datachain.query import Session
 
@@ -19,8 +16,8 @@ if TYPE_CHECKING:
 def read_pandas(  # type: ignore[override]
     df: "pd.DataFrame",
     name: str = "",
-    session: Optional[Session] = None,
-    settings: Optional[dict] = None,
+    session: Session | None = None,
+    settings: dict | None = None,
     in_memory: bool = False,
     column: str = "",
 ) -> "DataChain":

datachain/lib/dc/parquet.py

@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any
 
 from datachain.lib.data_model import DataType
 from datachain.query import Session
@@ -13,14 +13,14 @@ if TYPE_CHECKING:
 
 
 def read_parquet(
-    path: Union[str, os.PathLike[str], list[str], list[os.PathLike[str]]],
+    path: str | os.PathLike[str] | list[str] | list[os.PathLike[str]],
     partitioning: Any = "hive",
-    output: Optional[dict[str, DataType]] = None,
+    output: dict[str, DataType] | None = None,
     column: str = "",
     model_name: str = "",
     source: bool = True,
-    session: Optional[Session] = None,
-    settings: Optional[dict] = None,
+    session: Session | None = None,
+    settings: dict | None = None,
     **kwargs,
 ) -> "DataChain":
     """Generate chain from parquet files.

datachain/lib/dc/records.py

@@ -1,5 +1,5 @@
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 import sqlalchemy
 
@@ -19,11 +19,11 @@ READ_RECORDS_BATCH_SIZE = 10000
 
 
 def read_records(
-    to_insert: Optional[Union[dict, Iterable[dict]]],
-    session: Optional[Session] = None,
-    settings: Optional[dict] = None,
+    to_insert: dict | Iterable[dict] | None,
+    session: Session | None = None,
+    settings: dict | None = None,
     in_memory: bool = False,
-    schema: Optional[dict[str, DataType]] = None,
+    schema: dict[str, DataType] | None = None,
 ) -> "DataChain":
     """Create a DataChain from the provided records. This method can be used for
     programmatically generating a chain in contrast of reading data from storages

datachain/lib/dc/storage.py

@@ -1,7 +1,7 @@
 import os
 from collections.abc import Sequence
 from functools import reduce
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 from datachain.lib.dc.storage_pattern import (
     apply_glob_filter,
@@ -19,27 +19,27 @@ if TYPE_CHECKING:
 
 
 def read_storage(
-    uri: Union[str, os.PathLike[str], list[str], list[os.PathLike[str]]],
+    uri: str | os.PathLike[str] | list[str] | list[os.PathLike[str]],
     *,
     type: FileType = "binary",
-    session: Optional[Session] = None,
-    settings: Optional[dict] = None,
+    session: Session | None = None,
+    settings: dict | None = None,
     in_memory: bool = False,
-    recursive: Optional[bool] = True,
+    recursive: bool | None = True,
     column: str = "file",
     update: bool = False,
-    anon: Optional[bool] = None,
-    delta: Optional[bool] = False,
-    delta_on: Optional[Union[str, Sequence[str]]] = (
+    anon: bool | None = None,
+    delta: bool | None = False,
+    delta_on: str | Sequence[str] | None = (
         "file.path",
         "file.etag",
         "file.version",
     ),
-    delta_result_on: Optional[Union[str, Sequence[str]]] = None,
-    delta_compare: Optional[Union[str, Sequence[str]]] = None,
-    delta_retry: Optional[Union[bool, str]] = None,
+    delta_result_on: str | Sequence[str] | None = None,
+    delta_compare: str | Sequence[str] | None = None,
+    delta_retry: bool | str | None = None,
     delta_unsafe: bool = False,
-    client_config: Optional[dict] = None,
+    client_config: dict | None = None,
 ) -> "DataChain":
     """Get data from storage(s) as a list of file with all file attributes.
     It returns the chain itself as usual.

datachain/lib/dc/storage_pattern.py

@@ -1,5 +1,5 @@
 import glob
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 from datachain.client.fsspec import is_cloud_uri
 from datachain.lib.listing import ls
@@ -32,7 +32,7 @@ def validate_cloud_bucket_name(uri: str) -> None:
             raise ValueError(f"Glob patterns in bucket names are not supported: {uri}")
 
 
-def split_uri_pattern(uri: str) -> tuple[str, Union[str, None]]:
+def split_uri_pattern(uri: str) -> tuple[str, str | None]:
     """Split a URI into base path and glob pattern."""
     if not any(char in uri for char in ["*", "?", "[", "{", "}"]):
         return uri, None

datachain/lib/dc/utils.py

@@ -1,12 +1,6 @@
 from collections.abc import Sequence
 from functools import wraps
-from typing import (
-    TYPE_CHECKING,
-    Callable,
-    Optional,
-    TypeVar,
-    Union,
-)
+from typing import TYPE_CHECKING, TypeVar
 
 import sqlalchemy
 from sqlalchemy.sql.functions import GenericFunction
@@ -18,7 +12,10 @@ from datachain.query.schema import DEFAULT_DELIMITER
 from datachain.utils import getenv_bool
 
 if TYPE_CHECKING:
-    from typing_extensions import Concatenate, ParamSpec
+    from collections.abc import Callable
+    from typing import Concatenate
+
+    from typing_extensions import ParamSpec
 
     from .datachain import DataChain
 
@@ -70,11 +67,11 @@ class DatasetFromValuesError(DataChainParamsError):
         super().__init__(f"Dataset{name} from values error: {msg}")
 
 
-MergeColType = Union[str, Function, sqlalchemy.ColumnElement]
+MergeColType = str | Function | sqlalchemy.ColumnElement
 
 
 def _validate_merge_on(
-    on: Union[MergeColType, Sequence[MergeColType]],
+    on: MergeColType | Sequence[MergeColType],
     ds: "DataChain",
 ) -> Sequence[MergeColType]:
     if isinstance(on, (str, sqlalchemy.ColumnElement)):
@@ -103,12 +100,12 @@ def _get_merge_error_str(col: MergeColType) -> str:
 class DatasetMergeError(DataChainParamsError):
     def __init__(
         self,
-        on: Union[MergeColType, Sequence[MergeColType]],
-        right_on: Optional[Union[MergeColType, Sequence[MergeColType]]],
+        on: MergeColType | Sequence[MergeColType],
+        right_on: MergeColType | Sequence[MergeColType] | None,
         msg: str,
     ):
         def _get_str(
-            on: Union[MergeColType, Sequence[MergeColType]],
+            on: MergeColType | Sequence[MergeColType],
         ) -> str:
             if not isinstance(on, Sequence):
                 return str(on)  # type: ignore[unreachable]
@@ -123,7 +120,7 @@ class DatasetMergeError(DataChainParamsError):
         super().__init__(f"Merge error on='{on_str}'{right_on_str}: {msg}")
 
 
-OutputType = Union[None, DataType, Sequence[str], dict[str, DataType]]
+OutputType = DataType | Sequence[str] | dict[str, DataType] | None
 
 
 class Sys(DataModel):

datachain/lib/dc/values.py

@@ -1,8 +1,5 @@
 from collections.abc import Iterator
-from typing import (
-    TYPE_CHECKING,
-    Optional,
-)
+from typing import TYPE_CHECKING
 
 from datachain.lib.convert.values_to_tuples import values_to_tuples
 from datachain.lib.data_model import dict_to_data_model
@@ -20,8 +17,8 @@ if TYPE_CHECKING:
 
 def read_values(
     ds_name: str = "",
-    session: Optional[Session] = None,
-    settings: Optional[dict] = None,
+    session: Session | None = None,
+    settings: dict | None = None,
     in_memory: bool = False,
     output: OutputType = None,
     column: str = "",

datachain/lib/file.py

@@ -13,7 +13,7 @@ from datetime import datetime
 from functools import partial
 from io import BytesIO
 from pathlib import Path, PurePath, PurePosixPath
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Literal
 from urllib.parse import unquote, urlparse
 from urllib.request import url2pathname
 
@@ -53,12 +53,12 @@ class FileExporter(NodesThreadPool):
 
     def __init__(
         self,
-        output: Union[str, os.PathLike[str]],
+        output: str | os.PathLike[str],
         placement: ExportPlacement,
         use_cache: bool,
         link_type: Literal["copy", "symlink"],
         max_threads: int = EXPORT_FILES_MAX_THREADS,
-        client_config: Optional[dict] = None,
+        client_config: dict | None = None,
     ):
         super().__init__(max_threads)
         self.output = output
@@ -221,7 +221,7 @@ class File(DataModel):
     etag: str = Field(default="")
     is_latest: bool = Field(default=True)
     last_modified: datetime = Field(default=TIME_ZERO)
-    location: Optional[Union[dict, list[dict]]] = Field(default=None)
+    location: dict | list[dict] | None = Field(default=None)
 
     _datachain_column_types: ClassVar[dict[str, Any]] = {
         "source": String,
@@ -264,8 +264,8 @@ class File(DataModel):
 
     @staticmethod
     def _validate_dict(
-        v: Optional[Union[str, dict, list[dict]]],
-    ) -> Optional[Union[str, dict, list[dict]]]:
+        v: str | dict | list[dict] | None,
+    ) -> str | dict | list[dict] | None:
         if v is None or v == "":
             return None
         if isinstance(v, str):
@@ -334,8 +334,8 @@ class File(DataModel):
     def upload(
         cls,
         data: bytes,
-        path: Union[str, os.PathLike[str]],
-        catalog: Optional["Catalog"] = None,
+        path: str | os.PathLike[str],
+        catalog: "Catalog | None" = None,
     ) -> "Self":
         if catalog is None:
             from datachain.catalog.loader import get_catalog
@@ -357,7 +357,7 @@ class File(DataModel):
 
     @classmethod
     def at(
-        cls, uri: Union[str, os.PathLike[str]], session: Optional["Session"] = None
+        cls, uri: str | os.PathLike[str], session: "Session | None" = None
     ) -> "Self":
         """Construct a File from a full URI in one call.
 
@@ -470,7 +470,7 @@ class File(DataModel):
         """Returns file contents."""
         return self.read_bytes(length)
 
-    def save(self, destination: str, client_config: Optional[dict] = None):
+    def save(self, destination: str, client_config: dict | None = None):
         """Writes it's content to destination"""
         destination = stringify_path(destination)
         client: Client = self._catalog.get_client(destination, **(client_config or {}))
@@ -497,11 +497,11 @@ class File(DataModel):
 
     def export(
         self,
-        output: Union[str, os.PathLike[str]],
+        output: str | os.PathLike[str],
         placement: ExportPlacement = "fullpath",
         use_cache: bool = True,
         link_type: Literal["copy", "symlink"] = "copy",
-        client_config: Optional[dict] = None,
+        client_config: dict | None = None,
     ) -> None:
         """Export file to new location."""
         self._caching_enabled = use_cache
@@ -537,7 +537,7 @@ class File(DataModel):
         client = self._catalog.get_client(self.source)
         client.download(self, callback=self._download_cb)
 
-    async def _prefetch(self, download_cb: Optional["Callback"] = None) -> bool:
+    async def _prefetch(self, download_cb: "Callback | None" = None) -> bool:
         if self._catalog is None:
             raise RuntimeError("cannot prefetch file because catalog is not setup")
 
@@ -552,7 +552,7 @@ class File(DataModel):
         )
         return True
 
-    def get_local_path(self) -> Optional[str]:
+    def get_local_path(self) -> str | None:
         """Return path to a file in a local cache.
 
         Returns None if file is not cached.
@@ -629,7 +629,7 @@ class File(DataModel):
         return path
 
     def get_destination_path(
-        self, output: Union[str, os.PathLike[str]], placement: ExportPlacement
+        self, output: str | os.PathLike[str], placement: ExportPlacement
     ) -> str:
         """
         Returns full destination path of a file for exporting to some output
@@ -681,7 +681,7 @@ class File(DataModel):
             normalized_path = self.get_path_normalized()
             info = client.fs.info(client.get_full_path(normalized_path))
             converted_info = client.info_to_file(info, normalized_path)
-            return type(self)(
+            res = type(self)(
                 path=self.path,
                 source=self.source,
                 size=converted_info.size,
@@ -691,6 +691,8 @@ class File(DataModel):
                 last_modified=converted_info.last_modified,
                 location=self.location,
             )
+            res._set_stream(self._catalog)
+            return res
         except FileError as e:
             logger.warning(
                 "File error when resolving %s/%s: %s", self.source, self.path, str(e)
@@ -703,7 +705,7 @@ class File(DataModel):
                 str(e),
             )
 
-        return type(self)(
+        res = type(self)(
             path=self.path,
             source=self.source,
             size=0,
@@ -713,6 +715,8 @@ class File(DataModel):
             last_modified=TIME_ZERO,
             location=self.location,
         )
+        res._set_stream(self._catalog)
+        return res
 
     def rebase(
         self,
@@ -796,7 +800,7 @@ class TextFile(File):
         with self.open(**open_kwargs) as stream:
             return stream.read()
 
-    def save(self, destination: str, client_config: Optional[dict] = None):
+    def save(self, destination: str, client_config: dict | None = None):
         """Writes it's content to destination"""
         destination = stringify_path(destination)
 
@@ -829,8 +833,8 @@ class ImageFile(File):
     def save(  # type: ignore[override]
         self,
         destination: str,
-        format: Optional[str] = None,
-        client_config: Optional[dict] = None,
+        format: str | None = None,
+        client_config: dict | None = None,
     ):
         """Writes it's content to destination"""
         destination = stringify_path(destination)
@@ -912,7 +916,7 @@ class VideoFile(File):
     def get_frames(
         self,
         start: int = 0,
-        end: Optional[int] = None,
+        end: int | None = None,
         step: int = 1,
     ) -> "Iterator[VideoFrame]":
         """
@@ -962,7 +966,7 @@ class VideoFile(File):
         self,
         duration: float,
         start: float = 0,
-        end: Optional[float] = None,
+        end: float | None = None,
     ) -> "Iterator[VideoFragment]":
         """
         Splits the video into multiple fragments of a specified duration.
@@ -1048,7 +1052,7 @@ class AudioFile(File):
         self,
         duration: float,
         start: float = 0,
-        end: Optional[float] = None,
+        end: float | None = None,
     ) -> "Iterator[AudioFragment]":
         """
         Splits the audio into multiple fragments of a specified duration.
@@ -1086,10 +1090,10 @@ class AudioFile(File):
     def save(  # type: ignore[override]
         self,
         output: str,
-        format: Optional[str] = None,
+        format: str | None = None,
         start: float = 0,
-        end: Optional[float] = None,
-        client_config: Optional[dict] = None,
+        end: float | None = None,
+        client_config: dict | None = None,
     ) -> "AudioFile":
         """Save audio file or extract fragment to specified format.
 
@@ -1160,7 +1164,7 @@ class AudioFragment(DataModel):
         duration = self.end - self.start
         return audio_to_bytes(self.audio, format, self.start, duration)
 
-    def save(self, output: str, format: Optional[str] = None) -> "AudioFile":
+    def save(self, output: str, format: str | None = None) -> "AudioFile":
         """
         Saves the audio fragment as a new audio file.
 
@@ -1263,7 +1267,7 @@ class VideoFragment(DataModel):
     start: float
     end: float
 
-    def save(self, output: str, format: Optional[str] = None) -> "VideoFile":
+    def save(self, output: str, format: str | None = None) -> "VideoFile":
         """
         Saves the video fragment as a new video file.
 

datachain/lib/hf.py

@@ -26,7 +26,7 @@ except ImportError as exc:
     ) from exc
 
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, TypeAlias
 
 import PIL
 from tqdm.auto import tqdm
@@ -41,7 +41,9 @@ if TYPE_CHECKING:
     from pydantic import BaseModel
 
 
-HFDatasetType = Union[DatasetDict, Dataset, IterableDatasetDict, IterableDataset]
+HFDatasetType: TypeAlias = (
+    str | DatasetDict | Dataset | IterableDatasetDict | IterableDataset
+)
 
 
 class HFClassLabel(DataModel):
@@ -67,7 +69,7 @@ class HFAudio(DataModel):
 class HFGenerator(Generator):
     def __init__(
         self,
-        ds: Union[str, HFDatasetType],
+        ds: HFDatasetType,
         output_schema: type["BaseModel"],
         limit: int = 0,
         *args,
@@ -117,7 +119,7 @@ class HFGenerator(Generator):
                 pbar.update(1)
 
 
-def stream_splits(ds: Union[str, HFDatasetType], *args, **kwargs):
+def stream_splits(ds: HFDatasetType, *args, **kwargs):
     if isinstance(ds, str):
         ds = load_dataset(ds, *args, **kwargs)
     if isinstance(ds, (DatasetDict, IterableDatasetDict)):
@@ -153,7 +155,7 @@ def convert_feature(val: Any, feat: Any, anno: Any) -> Any:
 
 
 def get_output_schema(
-    features: Features, existing_column_names: Optional[list[str]] = None
+    features: Features, existing_column_names: list[str] | None = None
 ) -> tuple[dict[str, DataType], dict[str, str]]:
     """
     Generate UDF output schema from Hugging Face datasets features. It normalizes the

datachain/lib/image.py

@@ -1,4 +1,4 @@
-from typing import Callable, Optional, Union
+from collections.abc import Callable
 
 import torch
 from PIL import Image as PILImage
@@ -6,7 +6,7 @@ from PIL import Image as PILImage
 from datachain.lib.file import File, FileError, Image, ImageFile
 
 
-def image_info(file: Union[File, ImageFile]) -> Image:
+def image_info(file: File | ImageFile) -> Image:
     """
     Returns image file information.
 
@@ -31,11 +31,11 @@ def image_info(file: Union[File, ImageFile]) -> Image:
 def convert_image(
     img: PILImage.Image,
     mode: str = "RGB",
-    size: Optional[tuple[int, int]] = None,
-    transform: Optional[Callable] = None,
-    encoder: Optional[Callable] = None,
-    device: Optional[Union[str, torch.device]] = None,
-) -> Union[PILImage.Image, torch.Tensor]:
+    size: tuple[int, int] | None = None,
+    transform: Callable | None = None,
+    encoder: Callable | None = None,
+    device: str | torch.device | None = None,
+) -> PILImage.Image | torch.Tensor:
     """
     Resize, transform, and otherwise convert an image.
 
@@ -71,13 +71,13 @@ def convert_image(
 
 
 def convert_images(
-    images: Union[PILImage.Image, list[PILImage.Image]],
+    images: PILImage.Image | list[PILImage.Image],
     mode: str = "RGB",
-    size: Optional[tuple[int, int]] = None,
-    transform: Optional[Callable] = None,
-    encoder: Optional[Callable] = None,
-    device: Optional[Union[str, torch.device]] = None,
-) -> Union[list[PILImage.Image], torch.Tensor]:
+    size: tuple[int, int] | None = None,
+    transform: Callable | None = None,
+    encoder: Callable | None = None,
+    device: str | torch.device | None = None,
+) -> list[PILImage.Image] | torch.Tensor:
     """
     Resize, transform, and otherwise convert one or more images.
 

datachain/lib/listing.py

@@ -2,10 +2,10 @@ import glob
 import logging
 import os
 import posixpath
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from contextlib import contextmanager
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, TypeVar
 
 from fsspec.asyn import get_loop
 from sqlalchemy.sql.expression import true
@@ -73,7 +73,7 @@ def get_file_info(uri: str, cache, client_config=None) -> File:
 def ls(
     dc: D,
     path: str,
-    recursive: Optional[bool] = True,
+    recursive: bool | None = True,
     column="file",
 ) -> D:
     """
@@ -150,8 +150,8 @@ def _reraise_as_client_error() -> Iterator[None]:
 
 
 def get_listing(
-    uri: Union[str, os.PathLike[str]], session: "Session", update: bool = False
-) -> tuple[Optional[str], str, str, bool]:
+    uri: str | os.PathLike[str], session: "Session", update: bool = False
+) -> tuple[str | None, str, str, bool]:
     """Returns correct listing dataset name that must be used for saving listing
     operation. It takes into account existing listings and reusability of those.
     It also returns boolean saying if returned dataset name is reused / already

datachain/lib/listing_info.py

@@ -1,5 +1,4 @@
 from datetime import datetime, timedelta, timezone
-from typing import Optional
 
 from datachain.client import Client
 from datachain.lib.dataset_info import DatasetInfo
@@ -17,7 +16,7 @@ class ListingInfo(DatasetInfo):
         return uri
 
     @property
-    def expires(self) -> Optional[datetime]:
+    def expires(self) -> datetime | None:
         if not self.finished_at:
             return None
         return self.finished_at + timedelta(seconds=LISTING_TTL)

datachain/lib/meta_formats.py

@@ -2,9 +2,8 @@ import csv
 import json
 import tempfile
 import uuid
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from pathlib import Path
-from typing import Callable
 
 import jmespath as jsp
 from pydantic import BaseModel, ConfigDict, Field, ValidationError  # noqa: F401

datachain/lib/model_store.py

@@ -1,6 +1,6 @@
 import inspect
 import logging
-from typing import Any, ClassVar, Optional
+from typing import Any, ClassVar
 
 from pydantic import BaseModel
 
@@ -39,7 +39,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 +77,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

datachain/lib/namespaces.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 from datachain.error import (
     NamespaceCreateNotAllowedError,
     NamespaceDeleteNotAllowedError,
@@ -10,7 +8,7 @@ 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.
@@ -42,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.
@@ -61,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.
 
@@ -77,7 +75,7 @@ def ls(session: Optional[Session] = None) -> list[Namespace]:
     return Session.get(session).catalog.metastore.list_namespaces()
 
 
-def delete_namespace(name: str, session: Optional[Session] = None) -> None:
+def delete_namespace(name: str, session: Session | None = None) -> None:
     """
     Removes a namespace by name.