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

datachain/lib/file.py

@@ -1,7 +1,6 @@
 import errno
 import hashlib
 import io
-import json
 import logging
 import os
 import posixpath
@@ -13,7 +12,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
 
@@ -21,6 +20,7 @@ from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from fsspec.utils import stringify_path
 from pydantic import Field, field_validator
 
+from datachain import json
 from datachain.client.fileslice import FileSlice
 from datachain.lib.data_model import DataModel
 from datachain.lib.utils import DataChainError, rebase_path
@@ -35,13 +35,14 @@ if TYPE_CHECKING:
     from datachain.catalog import Catalog
     from datachain.client.fsspec import Client
     from datachain.dataset import RowDict
+    from datachain.query.session import Session
 
 sha256 = partial(hashlib.sha256, usedforsecurity=False)
 
 logger = logging.getLogger("datachain")
 
 # how to create file path when exporting
-ExportPlacement = Literal["filename", "etag", "fullpath", "checksum"]
+ExportPlacement = Literal["filename", "etag", "fullpath", "checksum", "filepath"]
 
 FileType = Literal["binary", "text", "image", "video", "audio"]
 EXPORT_FILES_MAX_THREADS = 5
@@ -52,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
@@ -220,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,
@@ -252,10 +253,19 @@ class File(DataModel):
         "last_modified",
     ]
 
+    # Allowed kwargs we forward to TextIOWrapper
+    _TEXT_WRAPPER_ALLOWED: ClassVar[tuple[str, ...]] = (
+        "encoding",
+        "errors",
+        "newline",
+        "line_buffering",
+        "write_through",
+    )
+
     @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):
@@ -287,6 +297,16 @@ class File(DataModel):
         super().__init__(**kwargs)
         self._catalog = None
         self._caching_enabled: bool = False
+        self._download_cb: Callback = DEFAULT_CALLBACK
+
+    def __getstate__(self):
+        state = super().__getstate__()
+        # Exclude _catalog from pickling - it contains SQLAlchemy engine and other
+        # non-picklable objects. The catalog will be re-set by _set_stream() on the
+        # worker side when needed.
+        state["__dict__"] = state["__dict__"].copy()
+        state["__dict__"]["_catalog"] = None
+        return state
 
     def as_text_file(self) -> "TextFile":
         """Convert the file to a `TextFile` object."""
@@ -322,17 +342,21 @@ class File(DataModel):
 
     @classmethod
     def upload(
-        cls, data: bytes, path: str, catalog: Optional["Catalog"] = None
+        cls,
+        data: bytes,
+        path: str | os.PathLike[str],
+        catalog: "Catalog | None" = None,
     ) -> "Self":
         if catalog is None:
-            from datachain.catalog.loader import get_catalog
-
-            catalog = get_catalog()
+            from datachain.query.session import Session
 
+            catalog = Session.get().catalog
         from datachain.client.fsspec import Client
 
-        client_cls = Client.get_implementation(path)
-        source, rel_path = client_cls.split_url(path)
+        path_str = stringify_path(path)
+
+        client_cls = Client.get_implementation(path_str)
+        source, rel_path = client_cls.split_url(path_str)
 
         client = catalog.get_client(client_cls.get_uri(source))
         file = client.upload(data, rel_path)
@@ -341,6 +365,35 @@ class File(DataModel):
         file._set_stream(catalog)
         return file
 
+    @classmethod
+    def at(
+        cls, uri: str | os.PathLike[str], session: "Session | None" = None
+    ) -> "Self":
+        """Construct a File from a full URI in one call.
+
+        Example:
+            file = File.at("s3://bucket/path/to/output.png")
+            with file.open("wb") as f: ...
+        """
+        from datachain.client.fsspec import Client
+        from datachain.query.session import Session
+
+        if session is None:
+            session = Session.get()
+        catalog = session.catalog
+        uri_str = stringify_path(uri)
+        if uri_str.endswith(("/", os.sep)):
+            raise ValueError(
+                f"File.at directory URL/path given (trailing slash), got: {uri_str}"
+            )
+        client_cls = Client.get_implementation(uri_str)
+        uri_str = client_cls.path_to_uri(uri_str)
+        source, rel_path = client_cls.split_url(uri_str)
+        source_uri = client_cls.get_uri(source)
+        file = cls(source=source_uri, path=rel_path)
+        file._set_stream(catalog)
+        return file
+
     @classmethod
     def _from_row(cls, row: "RowDict") -> "Self":
         return cls(**{key: row[key] for key in cls._datachain_column_types})
@@ -354,28 +407,93 @@ class File(DataModel):
         return str(PurePosixPath(self.path).parent)
 
     @contextmanager
-    def open(self, mode: Literal["rb", "r"] = "rb") -> Iterator[Any]:
-        """Open the file and return a file object."""
-        if self.location:
-            with VFileRegistry.open(self, self.location) as f:  # type: ignore[arg-type]
-                yield f
+    def open(
+        self,
+        mode: str = "rb",
+        *,
+        client_config: dict[str, Any] | None = None,
+        **open_kwargs,
+    ) -> Iterator[Any]:
+        """Open the file and return a file-like object.
+
+        Supports both read ("rb", "r") and write modes (e.g. "wb", "w", "ab").
+        When opened in a write mode, metadata is refreshed after closing.
+        """
+        writing = any(ch in mode for ch in "wax+")
+        if self.location and writing:
+            raise VFileError(
+                "Writing to virtual file is not supported",
+                self.source,
+                self.path,
+            )
 
-        else:
+        if self._catalog is None:
+            raise RuntimeError("Cannot open file: catalog is not set")
+
+        base_cfg = getattr(self._catalog, "client_config", {}) or {}
+        merged_cfg = {**base_cfg, **(client_config or {})}
+        client: Client = self._catalog.get_client(self.source, **merged_cfg)
+
+        if not writing:
+            if self.location:
+                with VFileRegistry.open(self, self.location) as f:  # type: ignore[arg-type]
+                    yield self._wrap_text(f, mode, open_kwargs)
+                return
             if self._caching_enabled:
                 self.ensure_cached()
-            client: Client = self._catalog.get_client(self.source)
             with client.open_object(
                 self, use_cache=self._caching_enabled, cb=self._download_cb
             ) as f:
-                yield io.TextIOWrapper(f) if mode == "r" else f
+                yield self._wrap_text(f, mode, open_kwargs)
+            return
+
+        # write path
+        full_path = client.get_full_path(self.get_path_normalized())
+        with client.fs.open(full_path, mode, **open_kwargs) as f:
+            yield self._wrap_text(f, mode, open_kwargs)
+
+        version_hint = self._extract_write_version(f)
+
+        # refresh metadata pinned to the version that was just written
+        refreshed = client.get_file_info(
+            self.get_path_normalized(), version_id=version_hint
+        )
+        for k, v in refreshed.model_dump().items():
+            setattr(self, k, v)
+
+    def _wrap_text(self, f: Any, mode: str, open_kwargs: dict[str, Any]) -> Any:
+        """Return stream possibly wrapped for text."""
+        if "b" in mode or isinstance(f, io.TextIOBase):
+            return f
+        filtered = {
+            k: open_kwargs[k] for k in self._TEXT_WRAPPER_ALLOWED if k in open_kwargs
+        }
+        return io.TextIOWrapper(f, **filtered)
+
+    def _extract_write_version(self, handle: Any) -> str | None:
+        """Best-effort extraction of object version after a write.
+
+        S3 (s3fs) and Azure (adlfs) populate version_id on the handle.
+        GCS (gcsfs) populates generation. Azure and GCS require upstream
+        fixes to be released.
+        """
+        for attr in ("version_id", "generation"):
+            if value := getattr(handle, attr, None):
+                return value
+        return None
 
     def read_bytes(self, length: int = -1):
         """Returns file contents as bytes."""
         with self.open() as stream:
             return stream.read(length)
 
-    def read_text(self):
-        """Returns file contents as text."""
+    def read_text(self, **open_kwargs):
+        """Return file contents decoded as text.
+
+        **open_kwargs : Any
+            Extra keyword arguments forwarded to ``open(mode="r", ...)``
+            (e.g. ``encoding="utf-8"``, ``errors="ignore"``)
+        """
         if self.location:
             raise VFileError(
                 "Reading text from virtual file is not supported",
@@ -383,14 +501,14 @@ class File(DataModel):
                 self.path,
             )
 
-        with self.open(mode="r") as stream:
+        with self.open(mode="r", **open_kwargs) as stream:
             return stream.read()
 
     def read(self, length: int = -1):
         """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 {}))
@@ -417,11 +535,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
@@ -457,7 +575,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")
 
@@ -472,7 +590,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.
@@ -549,7 +667,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
@@ -564,6 +682,8 @@ class File(DataModel):
             source = urlparse(self.source)
             if source.scheme and source.scheme != "file":
                 path = posixpath.join(source.netloc, path)
+        elif placement == "filepath":
+            path = unquote(self.get_path_normalized())
         elif placement == "checksum":
             raise NotImplementedError("Checksum placement not implemented yet")
         else:
@@ -601,7 +721,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,
@@ -611,6 +731,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)
@@ -623,7 +745,7 @@ class File(DataModel):
                 str(e),
             )
 
-        return type(self)(
+        res = type(self)(
             path=self.path,
             source=self.source,
             size=0,
@@ -633,6 +755,8 @@ class File(DataModel):
             last_modified=TIME_ZERO,
             location=self.location,
         )
+        res._set_stream(self._catalog)
+        return res
 
     def rebase(
         self,
@@ -701,17 +825,30 @@ class TextFile(File):
     """`DataModel` for reading text files."""
 
     @contextmanager
-    def open(self, mode: Literal["rb", "r"] = "r"):
-        """Open the file and return a file object (default to text mode)."""
-        with super().open(mode=mode) as stream:
+    def open(
+        self,
+        mode: str = "r",
+        *,
+        client_config: dict[str, Any] | None = None,
+        **open_kwargs,
+    ) -> Iterator[Any]:
+        """Open the file and return a file-like object.
+        Default to text mode"""
+        with super().open(
+            mode=mode, client_config=client_config, **open_kwargs
+        ) as stream:
             yield stream
 
-    def read_text(self):
-        """Returns file contents as text."""
-        with self.open() as stream:
+    def read_text(self, **open_kwargs):
+        """Return file contents as text.
+
+        **open_kwargs : Any
+            Extra keyword arguments forwarded to ``open()`` (e.g. encoding).
+        """
+        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)
 
@@ -744,8 +881,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)
@@ -827,7 +964,7 @@ class VideoFile(File):
     def get_frames(
         self,
         start: int = 0,
-        end: Optional[int] = None,
+        end: int | None = None,
         step: int = 1,
     ) -> "Iterator[VideoFrame]":
         """
@@ -877,7 +1014,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.
@@ -963,7 +1100,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.
@@ -1001,10 +1138,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.
 
@@ -1075,7 +1212,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.
 
@@ -1178,7 +1315,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

@@ -1,14 +1,13 @@
 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
 
+from datachain import json
 from datachain.lib.data_model import DataModel  # noqa: F401
 from datachain.lib.file import TextFile