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

datachain/catalog/catalog.py

@@ -25,7 +25,6 @@ from typing import (
 )
 from uuid import uuid4
 
-import requests
 import sqlalchemy as sa
 from sqlalchemy import Column
 from tqdm.auto import tqdm
@@ -54,7 +53,6 @@ from datachain.error import (
 from datachain.lib.listing import get_listing
 from datachain.node import DirType, Node, NodeWithPath
 from datachain.nodes_thread_pool import NodesThreadPool
-from datachain.remote.studio import StudioClient
 from datachain.sql.types import DateTime, SQLType
 from datachain.utils import DataChainDir
 
@@ -162,6 +160,8 @@ class DatasetRowsFetcher(NodesThreadPool):
         max_threads: int = PULL_DATASET_MAX_THREADS,
         progress_bar=None,
     ):
+        from datachain.remote.studio import StudioClient
+
         super().__init__(max_threads)
         self._check_dependencies()
         self.metastore = metastore
@@ -234,6 +234,8 @@ class DatasetRowsFetcher(NodesThreadPool):
         return df.drop("sys__id", axis=1)
 
     def get_parquet_content(self, url: str):
+        import requests
+
         while True:
             if self.should_check_for_status():
                 self.check_for_status()
@@ -1130,6 +1132,8 @@ class Catalog:
         raise DatasetNotFoundError(f"Dataset with version uuid {uuid} not found.")
 
     def get_remote_dataset(self, name: str) -> DatasetRecord:
+        from datachain.remote.studio import StudioClient
+
         studio_client = StudioClient()
 
         info_response = studio_client.dataset_info(name)
@@ -1164,8 +1168,27 @@ class Catalog:
 
         return direct_dependencies
 
-    def ls_datasets(self, include_listing: bool = False) -> Iterator[DatasetListRecord]:
-        datasets = self.metastore.list_datasets()
+    def ls_datasets(
+        self, include_listing: bool = False, studio: bool = False
+    ) -> Iterator[DatasetListRecord]:
+        from datachain.remote.studio import StudioClient
+
+        if studio:
+            client = StudioClient()
+            response = client.ls_datasets()
+            if not response.ok:
+                raise DataChainError(response.message)
+            if not response.data:
+                return
+
+            datasets: Iterator[DatasetListRecord] = (
+                DatasetListRecord.from_dict(d)
+                for d in response.data
+                if not d.get("name", "").startswith(QUERY_DATASET_PREFIX)
+            )
+        else:
+            datasets = self.metastore.list_datasets()
+
         for d in datasets:
             if not d.is_bucket_listing or include_listing:
                 yield d
@@ -1173,9 +1196,12 @@ class Catalog:
     def list_datasets_versions(
         self,
         include_listing: bool = False,
+        studio: bool = False,
     ) -> Iterator[tuple[DatasetListRecord, "DatasetListVersion", Optional["Job"]]]:
         """Iterate over all dataset versions with related jobs."""
-        datasets = list(self.ls_datasets(include_listing=include_listing))
+        datasets = list(
+            self.ls_datasets(include_listing=include_listing, studio=studio)
+        )
 
         # preselect dataset versions jobs from db to avoid multiple queries
         jobs_ids: set[str] = {
@@ -1345,6 +1371,8 @@ class Catalog:
         if cp and not output:
             raise ValueError("Please provide output directory for instantiation")
 
+        from datachain.remote.studio import StudioClient
+
         studio_client = StudioClient()
 
         try:

datachain/catalog/loader.py

@@ -1,19 +1,13 @@
 import os
 from importlib import import_module
-from typing import Any, Optional
-
-from datachain.catalog import Catalog
-from datachain.data_storage import (
-    AbstractMetastore,
-    AbstractWarehouse,
-)
-from datachain.data_storage.serializer import deserialize
-from datachain.data_storage.sqlite import (
-    SQLiteMetastore,
-    SQLiteWarehouse,
-)
+from typing import TYPE_CHECKING, Any, Optional
+
 from datachain.utils import get_envs_by_prefix
 
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+    from datachain.data_storage import AbstractMetastore, AbstractWarehouse
+
 METASTORE_SERIALIZED = "DATACHAIN__METASTORE"
 METASTORE_IMPORT_PATH = "DATACHAIN_METASTORE"
 METASTORE_ARG_PREFIX = "DATACHAIN_METASTORE_ARG_"
@@ -27,6 +21,9 @@ IN_MEMORY_ERROR_MESSAGE = "In-memory is only supported on SQLite"
 
 
 def get_metastore(in_memory: bool = False) -> "AbstractMetastore":
+    from datachain.data_storage import AbstractMetastore
+    from datachain.data_storage.serializer import deserialize
+
     metastore_serialized = os.environ.get(METASTORE_SERIALIZED)
     if metastore_serialized:
         metastore_obj = deserialize(metastore_serialized)
@@ -45,6 +42,8 @@ def get_metastore(in_memory: bool = False) -> "AbstractMetastore":
     }
 
     if not metastore_import_path:
+        from datachain.data_storage.sqlite import SQLiteMetastore
+
         metastore_args["in_memory"] = in_memory
         return SQLiteMetastore(**metastore_args)
     if in_memory:
@@ -62,6 +61,9 @@ def get_metastore(in_memory: bool = False) -> "AbstractMetastore":
 
 
 def get_warehouse(in_memory: bool = False) -> "AbstractWarehouse":
+    from datachain.data_storage import AbstractWarehouse
+    from datachain.data_storage.serializer import deserialize
+
     warehouse_serialized = os.environ.get(WAREHOUSE_SERIALIZED)
     if warehouse_serialized:
         warehouse_obj = deserialize(warehouse_serialized)
@@ -80,6 +82,8 @@ def get_warehouse(in_memory: bool = False) -> "AbstractWarehouse":
     }
 
     if not warehouse_import_path:
+        from datachain.data_storage.sqlite import SQLiteWarehouse
+
         warehouse_args["in_memory"] = in_memory
         return SQLiteWarehouse(**warehouse_args)
     if in_memory:
@@ -121,7 +125,7 @@ def get_distributed_class(**kwargs):
 
 def get_catalog(
     client_config: Optional[dict[str, Any]] = None, in_memory: bool = False
-) -> Catalog:
+) -> "Catalog":
     """
     Function that creates Catalog instance with appropriate metastore
     and warehouse classes. Metastore class can be provided with env variable
@@ -133,6 +137,8 @@ def get_catalog(
     and name of variable after, e.g. if it accepts team_id as kwargs
     we can provide DATACHAIN_METASTORE_ARG_TEAM_ID=12345 env variable.
     """
+    from datachain.catalog import Catalog
+
     return Catalog(
         metastore=get_metastore(in_memory=in_memory),
         warehouse=get_warehouse(in_memory=in_memory),

datachain/cli/__init__.py

@@ -6,7 +6,6 @@ from multiprocessing import freeze_support
 from typing import Optional
 
 from datachain.cli.utils import get_logging_level
-from datachain.telemetry import telemetry
 
 from .commands import (
     clear_cache,
@@ -70,6 +69,8 @@ def main(argv: Optional[list[str]] = None) -> int:
         error, return_code = handle_general_exception(exc, args, logging_level)
         return return_code
     finally:
+        from datachain.telemetry import telemetry
+
         telemetry.send_cli_call(args.command, error=error)
 
 

datachain/cli/parser/studio.py

@@ -63,19 +63,31 @@ def add_auth_parser(subparsers, parent_parser) -> None:
         default=False,
         help="Use code-based authentication without browser",
     )
+    login_parser.add_argument(
+        "--local",
+        action="store_true",
+        default=False,
+        help="Save the token in the local project config",
+    )
 
     auth_logout_help = "Log out from Studio"
     auth_logout_description = (
         "Remove the Studio authentication token from global config."
     )
 
-    auth_subparser.add_parser(
+    logout_parser = auth_subparser.add_parser(
         "logout",
         parents=[parent_parser],
         description=auth_logout_description,
         help=auth_logout_help,
         formatter_class=CustomHelpFormatter,
     )
+    logout_parser.add_argument(
+        "--local",
+        action="store_true",
+        default=False,
+        help="Remove the token from the local project config",
+    )
 
     auth_team_help = "Set default team for Studio operations"
     auth_team_description = "Set the default team for Studio operations."

datachain/client/fsspec.py

@@ -17,10 +17,10 @@ from typing import (
     ClassVar,
     NamedTuple,
     Optional,
+    Union,
 )
 from urllib.parse import urlparse
 
-from botocore.exceptions import ClientError
 from dvc_objects.fs.system import reflink
 from fsspec.asyn import get_loop, sync
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
@@ -28,7 +28,6 @@ from tqdm.auto import tqdm
 
 from datachain.cache import Cache
 from datachain.client.fileslice import FileWrapper
-from datachain.error import ClientError as DataChainClientError
 from datachain.nodes_fetcher import NodesFetcher
 from datachain.nodes_thread_pool import NodeChunk
 
@@ -83,19 +82,17 @@ class Client(ABC):
         self.uri = self.get_uri(self.name)
 
     @staticmethod
-    def get_implementation(url: str) -> type["Client"]:
+    def get_implementation(url: Union[str, os.PathLike[str]]) -> type["Client"]:
         from .azure import AzureClient
         from .gcs import GCSClient
         from .hf import HfClient
         from .local import FileClient
         from .s3 import ClientS3
 
-        protocol = urlparse(url).scheme
+        protocol = urlparse(str(url)).scheme
 
-        if not protocol or _is_win_local_path(url):
+        if not protocol or _is_win_local_path(str(url)):
             return FileClient
-
-        protocol = protocol.lower()
         if protocol == ClientS3.protocol:
             return ClientS3
         if protocol == GCSClient.protocol:
@@ -121,9 +118,11 @@ class Client(ABC):
         return cls.get_uri(storage_name), rel_path
 
     @staticmethod
-    def get_client(source: str, cache: Cache, **kwargs) -> "Client":
+    def get_client(
+        source: Union[str, os.PathLike[str]], cache: Cache, **kwargs
+    ) -> "Client":
         cls = Client.get_implementation(source)
-        storage_url, _ = cls.split_url(source)
+        storage_url, _ = cls.split_url(str(source))
         if os.name == "nt":
             storage_url = storage_url.removeprefix("/")
 
@@ -209,7 +208,7 @@ class Client(ABC):
 
     async def get_current_etag(self, file: "File") -> str:
         kwargs = {}
-        if self.fs.version_aware:
+        if getattr(self.fs, "version_aware", False):
             kwargs["version_id"] = file.version
         info = await self.fs._info(
             self.get_full_path(file.path, file.version), **kwargs
@@ -286,11 +285,6 @@ class Client(ABC):
                     worker.cancel()
             if excs:
                 raise excs[0]
-        except ClientError as exc:
-            raise DataChainClientError(
-                exc.response.get("Error", {}).get("Message") or exc,
-                exc.response.get("Error", {}).get("Code"),
-            ) from exc
         finally:
             # This ensures the progress bar is closed before any exceptions are raised
             progress_bar.close()
@@ -333,7 +327,9 @@ class Client(ABC):
         return not (key.startswith("/") or key.endswith("/") or "//" in key)
 
     async def ls_dir(self, path):
-        return await self.fs._ls(path, detail=True, versions=True)
+        if getattr(self.fs, "version_aware", False):
+            kwargs = {"versions": True}
+        return await self.fs._ls(path, detail=True, **kwargs)
 
     def rel_path(self, path: str) -> str:
         return self.fs.split_path(path)[1]

datachain/client/hf.py

@@ -1,25 +1,50 @@
-import os
+import functools
 import posixpath
-from typing import Any, cast
-
-from huggingface_hub import HfFileSystem
+from typing import Any
 
 from datachain.lib.file import File
 
 from .fsspec import Client
 
 
+class classproperty:  # noqa: N801
+    def __init__(self, func):
+        self.fget = func
+
+    def __get__(self, instance, owner):
+        return self.fget(owner)
+
+
+@functools.cache
+def get_hf_filesystem_cls():
+    import fsspec
+    from packaging.version import Version, parse
+
+    fsspec_version = parse(fsspec.__version__)
+    minver = Version("2024.12.0")
+
+    if fsspec_version < minver:
+        raise ImportError(
+            f"datachain requires 'fsspec>={minver}' but version "
+            f"{fsspec_version} is installed."
+        )
+
+    from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
+    from huggingface_hub import HfFileSystem
+
+    fs_cls = AsyncFileSystemWrapper.wrap_class(HfFileSystem)
+    # AsyncFileSystemWrapper does not set class properties, so we need to set them back.
+    fs_cls.protocol = HfFileSystem.protocol
+    return fs_cls
+
+
 class HfClient(Client):
-    FS_CLASS = HfFileSystem
     PREFIX = "hf://"
     protocol = "hf"
 
-    @classmethod
-    def create_fs(cls, **kwargs) -> HfFileSystem:
-        if os.environ.get("HF_TOKEN"):
-            kwargs["token"] = os.environ["HF_TOKEN"]
-
-        return cast(HfFileSystem, super().create_fs(**kwargs))
+    @classproperty
+    def FS_CLASS(cls):  # noqa: N802, N805
+        return get_hf_filesystem_cls()
 
     def info_to_file(self, v: dict[str, Any], path: str) -> File:
         return File(
@@ -31,8 +56,5 @@ class HfClient(Client):
             last_modified=v["last_commit"].date,
         )
 
-    async def ls_dir(self, path):
-        return self.fs.ls(path, detail=True)
-
     def rel_path(self, path):
         return posixpath.relpath(path, self.name)

datachain/client/local.py

@@ -67,10 +67,7 @@ class FileClient(Client):
     @classmethod
     def split_url(cls, url: str) -> tuple[str, str]:
         parsed = urlparse(url)
-        if parsed.scheme == "file":
-            scheme, rest = url.split(":", 1)
-            url = f"{scheme.lower()}:{rest}"
-        else:
+        if parsed.scheme != "file":
             url = cls.path_to_uri(url)
 
         fill_path = url[len(cls.PREFIX) :]

datachain/data_storage/warehouse.py

@@ -39,13 +39,6 @@ if TYPE_CHECKING:
     from datachain.data_storage.schema import DataTable
     from datachain.lib.file import File
 
-try:
-    import numpy as np
-
-    numpy_imported = True
-except ImportError:
-    numpy_imported = False
-
 
 logger = logging.getLogger("datachain")
 
@@ -96,7 +89,9 @@ class AbstractWarehouse(ABC, Serializable):
         If value is a list or some other iterable, it tries to convert sub elements
         as well
         """
-        if numpy_imported and isinstance(val, (np.ndarray, np.generic)):
+        import numpy as np
+
+        if isinstance(val, (np.ndarray, np.generic)):
             val = val.tolist()
 
         # Optimization: Precompute all the column type variables.

datachain/dataset.py

@@ -302,6 +302,7 @@ class DatasetListVersion:
         size: Optional[int],
         query_script: str = "",
         job_id: Optional[str] = None,
+        **kwargs,
     ):
         return cls(
             id,
@@ -648,6 +649,13 @@ class DatasetListRecord:
     def has_version_with_uuid(self, uuid: str) -> bool:
         return any(v.uuid == uuid for v in self.versions)
 
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "DatasetListRecord":
+        versions = [DatasetListVersion.parse(**v) for v in d.get("versions", [])]
+        kwargs = {f.name: d[f.name] for f in fields(cls) if f.name in d}
+        kwargs["versions"] = versions
+        return cls(**kwargs)
+
 
 class RowDict(dict):
     pass

datachain/error.py

@@ -1,15 +1,3 @@
-import botocore.errorfactory
-import botocore.exceptions
-import gcsfs.retry
-
-REMOTE_ERRORS = (
-    gcsfs.retry.HttpError,  # GCS
-    OSError,  # GCS
-    botocore.exceptions.BotoCoreError,  # S3
-    ValueError,  # Azure
-)
-
-
 class DataChainError(RuntimeError):
     pass
 

datachain/fs/utils.py

@@ -0,0 +1,30 @@
+from typing import TYPE_CHECKING
+
+from fsspec.implementations.local import LocalFileSystem
+
+if TYPE_CHECKING:
+    from fsspec import AbstractFileSystem
+
+
+def _isdir(fs: "AbstractFileSystem", path: str) -> bool:
+    info = fs.info(path)
+    return info["type"] == "directory" or (
+        info["size"] == 0 and info["type"] == "file" and info["name"].endswith("/")
+    )
+
+
+def isfile(fs: "AbstractFileSystem", path: str) -> bool:
+    """
+    Returns True if uri points to a file.
+
+    Supports special directories on object storages, e.g.:
+    Google creates a zero byte file with the same name as the directory with a trailing
+    slash at the end.
+    """
+    if isinstance(fs, LocalFileSystem):
+        return fs.isfile(path)
+
+    try:
+        return not _isdir(fs, path)
+    except FileNotFoundError:
+        return False

datachain/func/__init__.py

@@ -18,6 +18,7 @@ from .aggregate import (
 from .array import contains, cosine_distance, euclidean_distance, length, sip_hash_64
 from .conditional import and_, case, greatest, ifelse, isnone, least, or_
 from .numeric import bit_and, bit_hamming_distance, bit_or, bit_xor, int_hash_64
+from .path import file_ext, file_stem, name, parent
 from .random import rand
 from .string import byte_hamming_distance
 from .window import window
@@ -40,6 +41,8 @@ __all__ = [
     "count",
     "dense_rank",
     "euclidean_distance",
+    "file_ext",
+    "file_stem",
     "first",
     "greatest",
     "ifelse",
@@ -50,7 +53,9 @@ __all__ = [
     "literal",
     "max",
     "min",
+    "name",
     "or_",
+    "parent",
     "path",
     "rand",
     "random",

datachain/func/func.py

@@ -3,7 +3,6 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
 from sqlalchemy import BindParameter, Case, ColumnElement, Integer, cast, desc
-from sqlalchemy.ext.hybrid import Comparator
 from sqlalchemy.sql import func as sa_func
 
 from datachain.lib.convert.python_to_sql import python_to_sql
@@ -75,6 +74,8 @@ class Func(Function):
 
     @property
     def _db_cols(self) -> Sequence[ColT]:
+        from sqlalchemy.ext.hybrid import Comparator
+
         return (
             [
                 col

datachain/lib/dc.py

@@ -22,7 +22,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 +54,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 +213,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? "
@@ -408,7 +406,7 @@ class DataChain:
     @classmethod
     def from_storage(
         cls,
-        uri,
+        uri: Union[str, os.PathLike[str]],
         *,
         type: FileType = "binary",
         session: Optional[Session] = None,
@@ -550,6 +548,8 @@ class DataChain:
             )
             ```
         """
+        from datachain.telemetry import telemetry
+
         query = DatasetQuery(
             name=name,
             version=version,
@@ -573,7 +573,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 +610,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]
@@ -701,9 +701,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 +732,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
             )
         ]
 
@@ -1195,6 +1208,8 @@ class DataChain:
         )
         ```
         """
+        from sqlalchemy.sql.sqltypes import NullType
+
         primitives = (bool, str, int, float)
 
         for col_name, expr in kwargs.items():
@@ -2542,7 +2557,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",