datachain 0.11.11__py3-none-any.whl → 0.13.0__py3-none-any.whl

datachain/lib/dc.py

@@ -6,6 +6,7 @@ import sys
 from collections.abc import Iterator, Sequence
 from functools import wraps
 from typing import (
+    IO,
     TYPE_CHECKING,
     Any,
     BinaryIO,
@@ -22,7 +23,6 @@ import orjson
 import sqlalchemy
 from pydantic import BaseModel
 from sqlalchemy.sql.functions import GenericFunction
-from sqlalchemy.sql.sqltypes import NullType
 from tqdm import tqdm
 
 from datachain.dataset import DatasetRecord
@@ -55,7 +55,6 @@ from datachain.query import Session
 from datachain.query.dataset import DatasetQuery, PartitionByType
 from datachain.query.schema import DEFAULT_DELIMITER, Column, ColumnMeta
 from datachain.sql.functions import path as pathfunc
-from datachain.telemetry import telemetry
 from datachain.utils import batched_it, inside_notebook, row_to_nested_dict
 
 if TYPE_CHECKING:
@@ -215,7 +214,7 @@ class DataChain:
         from mistralai.client import MistralClient
         from mistralai.models.chat_completion import ChatMessage
 
-        from datachain.dc import DataChain, Column
+        from datachain import DataChain, Column
 
         PROMPT = (
             "Was this bot dialog successful? "
@@ -272,6 +271,18 @@ class DataChain:
         self._setup: dict = setup or {}
         self._sys = _sys
 
+    def __repr__(self) -> str:
+        """Return a string representation of the chain."""
+        classname = self.__class__.__name__
+        if not self._effective_signals_schema.values:
+            return f"Empty {classname}"
+
+        import io
+
+        file = io.StringIO()
+        self.print_schema(file=file)
+        return file.getvalue()
+
     @property
     def schema(self) -> dict[str, DataType]:
         """Get schema of the chain."""
@@ -325,9 +336,9 @@ class DataChain:
         """Return `self.union(other)`."""
         return self.union(other)
 
-    def print_schema(self) -> None:
+    def print_schema(self, file: Optional[IO] = None) -> None:
         """Print schema of the chain."""
-        self._effective_signals_schema.print_tree()
+        self._effective_signals_schema.print_tree(file=file)
 
     def clone(self) -> "Self":
         """Make a copy of the chain in a new table."""
@@ -408,7 +419,7 @@ class DataChain:
     @classmethod
     def from_storage(
         cls,
-        uri,
+        uri: Union[str, os.PathLike[str]],
         *,
         type: FileType = "binary",
         session: Optional[Session] = None,
@@ -550,6 +561,8 @@ class DataChain:
             )
             ```
         """
+        from datachain.telemetry import telemetry
+
         query = DatasetQuery(
             name=name,
             version=version,
@@ -573,7 +586,7 @@ class DataChain:
     @classmethod
     def from_json(
         cls,
-        path,
+        path: Union[str, os.PathLike[str]],
         type: FileType = "text",
         spec: Optional[DataType] = None,
         schema_from: Optional[str] = "auto",
@@ -610,7 +623,7 @@ class DataChain:
             ```
         """
         if schema_from == "auto":
-            schema_from = path
+            schema_from = str(path)
 
         def jmespath_to_name(s: str):
             name_end = re.search(r"\W", s).start() if re.search(r"\W", s) else len(s)  # type: ignore[union-attr]
@@ -629,7 +642,8 @@ class DataChain:
                 model_name=model_name,
                 jmespath=jmespath,
                 nrows=nrows,
-            )
+            ),
+            "params": {"file": File},
         }
         # disable prefetch if nrows is set
         settings = {"prefetch": 0} if nrows else {}
@@ -701,9 +715,22 @@ class DataChain:
         in_memory: bool = False,
         object_name: str = "dataset",
         include_listing: bool = False,
+        studio: bool = False,
     ) -> "DataChain":
         """Generate chain with list of registered datasets.
 
+        Args:
+            session: Optional session instance. If not provided, uses default session.
+            settings: Optional dictionary of settings to configure the chain.
+            in_memory: If True, creates an in-memory session. Defaults to False.
+            object_name: Name of the output object in the chain. Defaults to "dataset".
+            include_listing: If True, includes listing datasets. Defaults to False.
+            studio: If True, returns datasets from Studio only,
+                otherwise returns all local datasets. Defaults to False.
+
+        Returns:
+            DataChain: A new DataChain instance containing dataset information.
+
         Example:
             ```py
             from datachain import DataChain
@@ -719,7 +746,7 @@ class DataChain:
         datasets = [
             DatasetInfo.from_models(d, v, j)
             for d, v, j in catalog.list_datasets_versions(
-                include_listing=include_listing
+                include_listing=include_listing, studio=studio
             )
         ]
 
@@ -760,7 +787,12 @@ class DataChain:
         )
 
     def save(  # type: ignore[override]
-        self, name: Optional[str] = None, version: Optional[int] = None, **kwargs
+        self,
+        name: Optional[str] = None,
+        version: Optional[int] = None,
+        description: Optional[str] = None,
+        labels: Optional[list[str]] = None,
+        **kwargs,
     ) -> "Self":
         """Save to a Dataset. It returns the chain itself.
 
@@ -768,11 +800,18 @@ class DataChain:
             name : dataset name. Empty name saves to a temporary dataset that will be
                 removed after process ends. Temp dataset are useful for optimization.
             version : version of a dataset. Default - the last version that exist.
+            description : description of a dataset.
+            labels : labels of a dataset.
         """
         schema = self.signals_schema.clone_without_sys_signals().serialize()
         return self._evolve(
             query=self._query.save(
-                name=name, version=version, feature_schema=schema, **kwargs
+                name=name,
+                version=version,
+                description=description,
+                labels=labels,
+                feature_schema=schema,
+                **kwargs,
             )
         )
 
@@ -990,8 +1029,9 @@ class DataChain:
         func: Optional[Union[Callable, UDFObjT]],
         params: Union[None, str, Sequence[str]],
         output: OutputType,
-        signal_map,
+        signal_map: dict[str, Callable],
     ) -> UDFObjT:
+        is_batch = target_class.is_input_batched
         is_generator = target_class.is_output_batched
         name = self.name or ""
 
@@ -1002,7 +1042,9 @@ class DataChain:
         if self._sys:
             signals_schema = SignalSchema({"sys": Sys}) | signals_schema
 
-        params_schema = signals_schema.slice(sign.params, self._setup)
+        params_schema = signals_schema.slice(
+            sign.params, self._setup, is_batch=is_batch
+        )
 
         return target_class._create(sign, params_schema)
 
@@ -1195,6 +1237,8 @@ class DataChain:
         )
         ```
         """
+        from sqlalchemy.sql.sqltypes import NullType
+
         primitives = (bool, str, int, float)
 
         for col_name, expr in kwargs.items():
@@ -2542,7 +2586,7 @@ class DataChain:
 
     def to_storage(
         self,
-        output: str,
+        output: Union[str, os.PathLike[str]],
         signal: str = "file",
         placement: FileExportPlacement = "fullpath",
         link_type: Literal["copy", "symlink"] = "copy",

datachain/lib/file.py

@@ -18,7 +18,6 @@ from urllib.request import url2pathname
 
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from fsspec.utils import stringify_path
-from PIL import Image as PilImage
 from pydantic import Field, field_validator
 
 from datachain.client.fileslice import FileSlice
@@ -52,7 +51,7 @@ class FileExporter(NodesThreadPool):
 
     def __init__(
         self,
-        output: str,
+        output: Union[str, os.PathLike[str]],
         placement: ExportPlacement,
         use_cache: bool,
         link_type: Literal["copy", "symlink"],
@@ -194,7 +193,14 @@ class File(DataModel):
         "last_modified": DateTime,
         "location": JSON,
     }
-    _hidden_fields: ClassVar[list[str]] = ["version", "source"]
+    _hidden_fields: ClassVar[list[str]] = [
+        "source",
+        "version",
+        "etag",
+        "is_latest",
+        "last_modified",
+        "location",
+    ]
 
     _unique_id_keys: ClassVar[list[str]] = [
         "source",
@@ -243,6 +249,30 @@ class File(DataModel):
         self._catalog = None
         self._caching_enabled: bool = False
 
+    def as_text_file(self) -> "TextFile":
+        """Convert the file to a `TextFile` object."""
+        if isinstance(self, TextFile):
+            return self
+        file = TextFile(**self.model_dump())
+        file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
+        return file
+
+    def as_image_file(self) -> "ImageFile":
+        """Convert the file to a `ImageFile` object."""
+        if isinstance(self, ImageFile):
+            return self
+        file = ImageFile(**self.model_dump())
+        file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
+        return file
+
+    def as_video_file(self) -> "VideoFile":
+        """Convert the file to a `VideoFile` object."""
+        if isinstance(self, VideoFile):
+            return self
+        file = VideoFile(**self.model_dump())
+        file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
+        return file
+
     @classmethod
     def upload(
         cls, data: bytes, path: str, catalog: Optional["Catalog"] = None
@@ -292,20 +322,20 @@ class File(DataModel):
             ) as f:
                 yield io.TextIOWrapper(f) if mode == "r" else f
 
-    def read(self, length: int = -1):
-        """Returns file contents."""
+    def read_bytes(self, length: int = -1):
+        """Returns file contents as bytes."""
         with self.open() as stream:
             return stream.read(length)
 
-    def read_bytes(self):
-        """Returns file contents as bytes."""
-        return self.read()
-
     def read_text(self):
         """Returns file contents as text."""
         with self.open(mode="r") 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):
         """Writes it's content to destination"""
         destination = stringify_path(destination)
@@ -333,7 +363,7 @@ class File(DataModel):
 
     def export(
         self,
-        output: str,
+        output: Union[str, os.PathLike[str]],
         placement: ExportPlacement = "fullpath",
         use_cache: bool = True,
         link_type: Literal["copy", "symlink"] = "copy",
@@ -374,15 +404,10 @@ class File(DataModel):
         client.download(self, callback=self._download_cb)
 
     async def _prefetch(self, download_cb: Optional["Callback"] = None) -> bool:
-        from datachain.client.hf import HfClient
-
         if self._catalog is None:
             raise RuntimeError("cannot prefetch file because catalog is not setup")
 
         client = self._catalog.get_client(self.source)
-        if client.protocol == HfClient.protocol:
-            return False
-
         await client._download(self, callback=download_cb or self._download_cb)
         self._set_stream(
             self._catalog, caching_enabled=True, download_cb=DEFAULT_CALLBACK
@@ -430,7 +455,9 @@ class File(DataModel):
             path = url2pathname(path)
         return path
 
-    def get_destination_path(self, output: str, placement: ExportPlacement) -> str:
+    def get_destination_path(
+        self, output: Union[str, os.PathLike[str]], placement: ExportPlacement
+    ) -> str:
         """
         Returns full destination path of a file for exporting to some output
         based on export placement
@@ -551,18 +578,36 @@ class TextFile(File):
 class ImageFile(File):
     """`DataModel` for reading image files."""
 
+    def get_info(self) -> "Image":
+        """
+        Retrieves metadata and information about the image file.
+
+        Returns:
+            Image: A Model containing image metadata such as width, height and format.
+        """
+        from .image import image_info
+
+        return image_info(self)
+
     def read(self):
         """Returns `PIL.Image.Image` object."""
+        from PIL import Image as PilImage
+
         fobj = super().read()
         return PilImage.open(BytesIO(fobj))
 
-    def save(self, destination: str, client_config: Optional[dict] = None):
+    def save(  # type: ignore[override]
+        self,
+        destination: str,
+        format: Optional[str] = None,
+        client_config: Optional[dict] = None,
+    ):
         """Writes it's content to destination"""
         destination = stringify_path(destination)
 
         client: Client = self._catalog.get_client(destination, **(client_config or {}))
         with client.fs.open(destination, mode="wb") as f:
-            self.read().save(f)
+            self.read().save(f, format=format)
 
 
 class Image(DataModel):

datachain/lib/image.py

@@ -1,17 +1,41 @@
 from typing import Callable, Optional, Union
 
 import torch
-from PIL import Image
+from PIL import Image as PILImage
+
+from datachain.lib.file import File, FileError, Image, ImageFile
+
+
+def image_info(file: Union[File, ImageFile]) -> Image:
+    """
+    Returns image file information.
+
+    Args:
+        file (ImageFile): Image file object.
+
+    Returns:
+        Image: Image file information.
+    """
+    try:
+        img = file.as_image_file().read()
+    except Exception as exc:
+        raise FileError(file, "unable to open image file") from exc
+
+    return Image(
+        width=img.width,
+        height=img.height,
+        format=img.format or "",
+    )
 
 
 def convert_image(
-    img: Image.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[Image.Image, torch.Tensor]:
+) -> Union[PILImage.Image, torch.Tensor]:
     """
     Resize, transform, and otherwise convert an image.
 
@@ -47,13 +71,13 @@ def convert_image(
 
 
 def convert_images(
-    images: Union[Image.Image, list[Image.Image]],
+    images: Union[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[Image.Image], torch.Tensor]:
+) -> Union[list[PILImage.Image], torch.Tensor]:
     """
     Resize, transform, and otherwise convert one or more images.
 
@@ -65,7 +89,7 @@ def convert_images(
         encoder (Callable): Encode image using model.
         device (str or torch.device): Device to use.
     """
-    if isinstance(images, Image.Image):
+    if isinstance(images, PILImage.Image):
         images = [images]
 
     converted = [

datachain/lib/listing.py

@@ -1,19 +1,21 @@
+import glob
 import logging
 import os
 import posixpath
 from collections.abc import Iterator
-from typing import TYPE_CHECKING, Callable, Optional, TypeVar
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union
 
 from fsspec.asyn import get_loop
 from sqlalchemy.sql.expression import true
 
+import datachain.fs.utils as fsutils
 from datachain.asyn import iter_over_async
 from datachain.client import Client
-from datachain.error import REMOTE_ERRORS, ClientError
+from datachain.error import ClientError
 from datachain.lib.file import File
 from datachain.query.schema import Column
 from datachain.sql.functions import path as pathfunc
-from datachain.telemetry import telemetry
 from datachain.utils import uses_glob
 
 if TYPE_CHECKING:
@@ -92,38 +94,6 @@ def ls(
     return dc.filter(pathfunc.parent(_file_c("path")) == path.lstrip("/").rstrip("/*"))
 
 
-def _isfile(client: "Client", path: str) -> bool:
-    """
-    Returns True if uri points to a file
-    """
-    try:
-        if "://" in path:
-            # This makes sure that the uppercase scheme is converted to lowercase
-            scheme, path = path.split("://", 1)
-            path = f"{scheme.lower()}://{path}"
-
-        if os.name == "nt" and "*" in path:
-            # On Windows, the glob pattern "*" is not supported
-            return False
-
-        info = client.fs.info(path)
-        name = info.get("name")
-        # case for special simulated directories on some clouds
-        # e.g. Google creates a zero byte file with the same name as the
-        # directory with a trailing slash at the end
-        if not name or name.endswith("/"):
-            return False
-
-        return info["type"] == "file"
-    except FileNotFoundError:
-        return False
-    except REMOTE_ERRORS as e:
-        raise ClientError(
-            message=str(e),
-            error_code=getattr(e, "code", None),
-        ) from e
-
-
 def parse_listing_uri(uri: str, client_config) -> tuple[str, str, str]:
     """
     Parsing uri and returns listing dataset name, listing uri and listing path
@@ -156,8 +126,16 @@ def listing_uri_from_name(dataset_name: str) -> str:
     return dataset_name.removeprefix(LISTING_PREFIX)
 
 
+@contextmanager
+def _reraise_as_client_error() -> Iterator[None]:
+    try:
+        yield
+    except Exception as e:
+        raise ClientError(message=str(e), error_code=getattr(e, "code", None)) from e
+
+
 def get_listing(
-    uri: str, session: "Session", update: bool = False
+    uri: Union[str, os.PathLike[str]], session: "Session", update: bool = False
 ) -> tuple[Optional[str], 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.
@@ -167,6 +145,7 @@ def get_listing(
     be used to find rows based on uri.
     """
     from datachain.client.local import FileClient
+    from datachain.telemetry import telemetry
 
     catalog = session.catalog
     cache = catalog.cache
@@ -174,11 +153,14 @@ def get_listing(
 
     client = Client.get_client(uri, cache, **client_config)
     telemetry.log_param("client", client.PREFIX)
+    if not isinstance(uri, str):
+        uri = os.fspath(uri)
 
     # we don't want to use cached dataset (e.g. for a single file listing)
-    if not uri.endswith("/") and _isfile(client, uri):
-        storage_uri, path = Client.parse_url(uri)
-        return None, f"{storage_uri}/{path.lstrip('/')}", path, False
+    isfile = _reraise_as_client_error()(fsutils.isfile)
+    if not glob.has_magic(uri) and not uri.endswith("/") and isfile(client.fs, uri):
+        _, path = Client.parse_url(uri)
+        return None, uri, path, False
 
     ds_name, list_uri, list_path = parse_listing_uri(uri, client_config)
     listing = None

datachain/lib/meta_formats.py

@@ -10,7 +10,7 @@ import jmespath as jsp
 from pydantic import BaseModel, ConfigDict, Field, ValidationError  # noqa: F401
 
 from datachain.lib.data_model import DataModel  # noqa: F401
-from datachain.lib.file import File
+from datachain.lib.file import TextFile
 
 
 class UserModel(BaseModel):
@@ -130,7 +130,7 @@ def read_meta(  # noqa: C901
     #
 
     def parse_data(
-        file: File,
+        file: TextFile,
         data_model=spec,
         format=format,
         jmespath=jmespath,

datachain/lib/signal_schema.py

@@ -5,6 +5,7 @@ from dataclasses import dataclass
 from datetime import datetime
 from inspect import isclass
 from typing import (  # noqa: UP035
+    IO,
     TYPE_CHECKING,
     Annotated,
     Any,
@@ -154,9 +155,9 @@ class SignalSchema:
             if not callable(func):
                 raise SetupError(key, "value must be function or callable class")
 
-    def _init_setup_values(self):
+    def _init_setup_values(self) -> None:
         if self.setup_values is not None:
-            return self.setup_values
+            return
 
         res = {}
         for key, func in self.setup_func.items():
@@ -398,7 +399,7 @@ class SignalSchema:
         return SignalSchema(signals)
 
     @staticmethod
-    def get_flatten_hidden_fields(schema):
+    def get_flatten_hidden_fields(schema: dict):
         custom_types = schema.get("_custom_types", {})
         if not custom_types:
             return []
@@ -464,19 +465,61 @@ class SignalSchema:
         return False
 
     def slice(
-        self, keys: Sequence[str], setup: Optional[dict[str, Callable]] = None
+        self,
+        params: dict[str, Union[DataType, Any]],
+        setup: Optional[dict[str, Callable]] = None,
+        is_batch: bool = False,
     ) -> "SignalSchema":
-        # Make new schema that combines current schema and setup signals
-        setup = setup or {}
-        setup_no_types = dict.fromkeys(setup.keys(), str)
-        union = SignalSchema(self.values | setup_no_types)
-        # Slice combined schema by keys
-        schema = {}
-        for k in keys:
-            try:
-                schema[k] = union._find_in_tree(k.split("."))
-            except SignalResolvingError:
-                pass
+        """
+        Returns new schema that combines current schema and setup signals.
+        """
+        setup_params = setup.keys() if setup else []
+        schema: dict[str, DataType] = {}
+
+        for param, param_type in params.items():
+            # This is special case for setup params, they are always treated as strings
+            if param in setup_params:
+                schema[param] = str
+                continue
+
+            schema_type = self._find_in_tree(param.split("."))
+
+            if param_type is Any:
+                schema[param] = schema_type
+                continue
+
+            schema_origin = get_origin(schema_type)
+            param_origin = get_origin(param_type)
+
+            if schema_origin is Union and type(None) in get_args(schema_type):
+                schema_type = get_args(schema_type)[0]
+                if param_origin is Union and type(None) in get_args(param_type):
+                    param_type = get_args(param_type)[0]
+
+            if is_batch:
+                if param_type is list:
+                    schema[param] = schema_type
+                    continue
+
+                if param_origin is not list:
+                    raise SignalResolvingError(param.split("."), "is not a list")
+
+                param_type = get_args(param_type)[0]
+
+            if param_type == schema_type or (
+                isclass(param_type)
+                and isclass(schema_type)
+                and issubclass(param_type, File)
+                and issubclass(schema_type, File)
+            ):
+                schema[param] = schema_type
+                continue
+
+            raise SignalResolvingError(
+                param.split("."),
+                f"types mismatch: {param_type} != {schema_type}",
+            )
+
         return SignalSchema(schema, setup)
 
     def row_to_features(
@@ -696,16 +739,20 @@ class SignalSchema:
                     substree, new_prefix, depth + 1, include_hidden
                 )
 
-    def print_tree(self, indent: int = 4, start_at: int = 0):
+    def print_tree(self, indent: int = 2, start_at: int = 0, file: Optional[IO] = None):
         for path, type_, _, depth in self.get_flat_tree():
             total_indent = start_at + depth * indent
-            print(" " * total_indent, f"{path[-1]}:", SignalSchema._type_to_str(type_))
+            col_name = " " * total_indent + path[-1]
+            col_type = SignalSchema._type_to_str(type_)
+            print(col_name, col_type, sep=": ", file=file)
 
             if get_origin(type_) is list:
                 args = get_args(type_)
                 if len(args) > 0 and ModelStore.is_pydantic(args[0]):
                     sub_schema = SignalSchema({"* list of": args[0]})
-                    sub_schema.print_tree(indent=indent, start_at=total_indent + indent)
+                    sub_schema.print_tree(
+                        indent=indent, start_at=total_indent + indent, file=file
+                    )
 
     def get_headers_with_length(self, include_hidden: bool = True):
         paths = [