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

datachain/cli/commands/ls.py

@@ -1,13 +1,14 @@
 import shlex
 from collections.abc import Iterable, Iterator
 from itertools import chain
-from typing import TYPE_CHECKING, Optional
-
-if TYPE_CHECKING:
-    from datachain.catalog import Catalog
+from typing import TYPE_CHECKING
 
 from datachain.cli.utils import determine_flavors
 from datachain.config import Config
+from datachain.query.session import Session
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
 
 
 def ls(
@@ -16,7 +17,7 @@ def ls(
     studio: bool = False,
     local: bool = False,
     all: bool = True,
-    team: Optional[str] = None,
+    team: str | None = None,
     **kwargs,
 ):
     token = Config().read().get("studio", {}).get("token")
@@ -32,18 +33,15 @@ def ls(
 def ls_local(
     sources,
     long: bool = False,
-    catalog: Optional["Catalog"] = None,
+    catalog=None,
     client_config=None,
     **kwargs,
 ):
     from datachain import listings
 
     if sources:
-        if catalog is None:
-            from datachain.catalog import get_catalog
-
-            catalog = get_catalog(client_config=client_config)
-
+        session = Session.get(catalog=catalog, client_config=client_config)
+        catalog = session.catalog
         actual_sources = list(ls_urls(sources, catalog=catalog, long=long, **kwargs))
         if len(actual_sources) == 1:
             for _, entries in actual_sources:
@@ -63,8 +61,8 @@ def ls_local(
                     print(format_ls_entry(entry))
     else:
         # Collect results in a list here to prevent interference from `tqdm` and `print`
-        listing = list(listings().collect("listing"))
-        for ls in listing:
+        listing = listings().to_list("listing")
+        for (ls,) in listing:
             print(format_ls_entry(f"{ls.uri}@v{ls.version}"))  # type: ignore[union-attr]
 
 
@@ -78,7 +76,7 @@ def format_ls_entry(entry: str) -> str:
 def ls_remote(
     paths: Iterable[str],
     long: bool = False,
-    team: Optional[str] = None,
+    team: str | None = None,
 ):
     from datachain.node import long_line_str
     from datachain.remote.studio import StudioClient
@@ -145,7 +143,7 @@ def _ls_urls_flat(
     long: bool,
     catalog: "Catalog",
     **kwargs,
-) -> Iterator[tuple[str, Iterator[str]]]:
+) -> Iterator[tuple[str, Iterable[str]]]:
     from datachain.client import Client
     from datachain.node import long_line_str
 
@@ -154,7 +152,9 @@ def _ls_urls_flat(
         if client_cls.is_root_url(source):
             buckets = client_cls.ls_buckets(**catalog.client_config)
             if long:
-                values = (long_line_str(b.name, b.created) for b in buckets)
+                values: Iterable[str] = (
+                    long_line_str(b.name, b.created) for b in buckets
+                )
             else:
                 values = (b.name for b in buckets)
             yield source, values
@@ -164,7 +164,7 @@ def _ls_urls_flat(
             if long:
                 fields.append("last_modified")
             for data_source, results in catalog.ls([source], fields=fields, **kwargs):
-                values = (_node_data_to_ls_values(r, long) for r in results)
+                values = [_node_data_to_ls_values(r, long) for r in results]
                 found = True
                 yield data_source.dirname(), values
             if not found:

datachain/cli/commands/show.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 from datachain.lib.signal_schema import SignalSchema
 
@@ -10,7 +10,7 @@ if TYPE_CHECKING:
 def show(
     catalog: "Catalog",
     name: str,
-    version: Optional[int] = None,
+    version: str | None = None,
     limit: int = 10,
     offset: int = 0,
     columns: Sequence[str] = (),
@@ -42,8 +42,8 @@ def show(
     print("Name: ", name)
     if dataset.description:
         print("Description: ", dataset.description)
-    if dataset.labels:
-        print("Labels: ", ",".join(dataset.labels))
+    if dataset.attrs:
+        print("Attributes: ", ",".join(dataset.attrs))
     print("\n")
 
     show_records(records, collapse_columns=not no_collapse, hidden_fields=hidden_fields)

datachain/cli/parser/__init__.py

@@ -3,7 +3,7 @@ from importlib.metadata import PackageNotFoundError, version
 
 import shtab
 
-from datachain.cli.utils import BooleanOptionalAction, KeyValueArgs
+from datachain.cli.utils import BooleanOptionalAction
 
 from .job import add_jobs_parser
 from .studio import add_auth_parser
@@ -16,9 +16,7 @@ from .utils import (
     add_update_arg,
     find_columns_type,
 )
-from .utils import (
-    CustomArgumentParser as ArgumentParser,
-)
+from .utils import CustomArgumentParser as ArgumentParser
 
 
 def get_parser() -> ArgumentParser:  # noqa: PLR0915
@@ -217,29 +215,9 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="Dataset description",
     )
     parse_edit_dataset.add_argument(
-        "--labels",
+        "--attrs",
         nargs="+",
-        help="Dataset labels",
-    )
-    parse_edit_dataset.add_argument(
-        "--studio",
-        action="store_true",
-        default=False,
-        help="Edit dataset from Studio",
-    )
-    parse_edit_dataset.add_argument(
-        "-L",
-        "--local",
-        action="store_true",
-        default=False,
-        help="Edit local dataset only",
-    )
-    parse_edit_dataset.add_argument(
-        "-a",
-        "--all",
-        action="store_true",
-        default=True,
-        help="Edit both datasets from studio and local",
+        help="Dataset attributes",
     )
     parse_edit_dataset.add_argument(
         "--team",
@@ -302,7 +280,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         "--version",
         action="store",
         default=None,
-        type=int,
+        type=str,
         help="Dataset version",
     )
     rm_dataset_parser.add_argument(
@@ -315,21 +293,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         "--studio",
         action="store_true",
         default=False,
-        help="Remove dataset from Studio",
-    )
-    rm_dataset_parser.add_argument(
-        "-L",
-        "--local",
-        action="store_true",
-        default=False,
-        help="Remove local datasets only",
-    )
-    rm_dataset_parser.add_argument(
-        "-a",
-        "--all",
-        action="store_true",
-        default=True,
-        help="Remove both local and studio",
+        help="Remove dataset from Studio only",
     )
     rm_dataset_parser.add_argument(
         "--team",
@@ -495,43 +459,12 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         "--version",
         action="store",
         default=None,
-        type=int,
+        type=str,
         help="Dataset version",
     )
     show_parser.add_argument("--schema", action="store_true", help="Show schema")
     add_show_args(show_parser)
 
-    query_parser = subp.add_parser(
-        "query",
-        parents=[parent_parser],
-        description="Create a new dataset with a query script.",
-        formatter_class=CustomHelpFormatter,
-    )
-    add_anon_arg(query_parser)
-    query_parser.add_argument(
-        "script", metavar="<script.py>", type=str, help="Filepath for script"
-    )
-    query_parser.add_argument(
-        "--parallel",
-        nargs="?",
-        type=int,
-        const=-1,
-        default=None,
-        metavar="N",
-        help=(
-            "Use multiprocessing to run any query script UDFs with N worker processes. "
-            "N defaults to the CPU count"
-        ),
-    )
-    query_parser.add_argument(
-        "-p",
-        "--param",
-        metavar="param=value",
-        nargs=1,
-        action=KeyValueArgs,
-        help="Query parameters",
-    )
-
     parse_clear_cache = subp.add_parser(
         "clear-cache",
         parents=[parent_parser],
@@ -550,6 +483,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
 
     subp.add_parser("internal-run-udf", parents=[parent_parser])
     subp.add_parser("internal-run-udf-worker", parents=[parent_parser])
+
     add_completion_parser(subp, [parent_parser])
     return parser
 

datachain/cli/parser/job.py

@@ -13,11 +13,16 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
     )
     jobs_subparser = jobs_parser.add_subparsers(
         dest="cmd",
-        help="Use `datachain auth CMD --help` to display command-specific help",
+        help="Use `datachain job CMD --help` to display command-specific help",
     )
 
     studio_run_help = "Run a job in Studio"
-    studio_run_description = "Run a job in Studio."
+    studio_run_description = "Run a job in Studio. \n"
+    studio_run_description += (
+        "When using --start-time or --cron,"
+        " the job is scheduled to run but won't start immediately"
+        " (can be seen in the Tasks tab in UI)"
+    )
 
     studio_run_parser = jobs_subparser.add_parser(
         "run",
@@ -51,6 +56,20 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         help="Environment variables in KEY=VALUE format",
     )
 
+    studio_run_parser.add_argument(
+        "--cluster",
+        type=str,
+        action="store",
+        help="Compute cluster to run the job on",
+    )
+
+    studio_run_parser.add_argument(
+        "-c",
+        "--credentials-name",
+        action="store",
+        help="Name of the credentials to use for the job",
+    )
+
     studio_run_parser.add_argument(
         "--workers",
         type=int,
@@ -64,7 +83,12 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
     studio_run_parser.add_argument(
         "--python-version",
         action="store",
-        help="Python version for the job (e.g., 3.9, 3.10, 3.11)",
+        help="Python version for the job (e.g., 3.10, 3.11, 3.12, 3.13)",
+    )
+    studio_run_parser.add_argument(
+        "--repository",
+        action="store",
+        help="Repository URL to clone before running the job",
     )
     studio_run_parser.add_argument(
         "--req-file",
@@ -77,6 +101,56 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         nargs="+",
         help="Python package requirements",
     )
+    studio_run_parser.add_argument(
+        "--priority",
+        type=int,
+        default=5,
+        help="Priority for the job in range 0-5. "
+        "Lower value is higher priority (default: 5)",
+    )
+    studio_run_parser.add_argument(
+        "--start-time",
+        action="store",
+        help="Time to schedule a task in YYYY-MM-DDTHH:mm format or natural language.",
+    )
+    studio_run_parser.add_argument(
+        "--cron", action="store", help="Cron expression for the cron task."
+    )
+    studio_run_parser.add_argument(
+        "--no-wait",
+        action="store_true",
+        help="Do not wait for the job to finish",
+    )
+
+    studio_ls_help = "List jobs in Studio"
+    studio_ls_description = "List jobs in Studio."
+
+    studio_ls_parser = jobs_subparser.add_parser(
+        "ls",
+        parents=[parent_parser],
+        description=studio_ls_description,
+        help=studio_ls_help,
+        formatter_class=CustomHelpFormatter,
+    )
+
+    studio_ls_parser.add_argument(
+        "--status",
+        action="store",
+        help="Status to filter jobs by",
+    )
+
+    studio_ls_parser.add_argument(
+        "--team",
+        action="store",
+        default=None,
+        help="Team to list jobs for (default: from config)",
+    )
+    studio_ls_parser.add_argument(
+        "--limit",
+        type=int,
+        default=20,
+        help="Limit the number of jobs returned (default: 20)",
+    )
 
     studio_cancel_help = "Cancel a job in Studio"
     studio_cancel_description = "Cancel a running job in Studio."
@@ -123,3 +197,21 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         default=None,
         help="Team to check logs for (default: from config)",
     )
+
+    studio_clusters_help = "List compute clusters in Studio"
+    studio_clusters_description = "List compute clusters in Studio."
+
+    studio_clusters_parser = jobs_subparser.add_parser(
+        "clusters",
+        parents=[parent_parser],
+        description=studio_clusters_description,
+        help=studio_clusters_help,
+        formatter_class=CustomHelpFormatter,
+    )
+
+    studio_clusters_parser.add_argument(
+        "--team",
+        action="store",
+        default=None,
+        help="Team to list clusters for (default: from config)",
+    )

datachain/cli/parser/studio.py

@@ -89,8 +89,13 @@ def add_auth_parser(subparsers, parent_parser) -> None:
         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."
+    auth_team_help = "Set or show default team for Studio operations"
+    auth_team_description = (
+        "Set or show the default team for Studio operations. "
+        "This will be used globally by default. "
+        "Use --local to set the team locally for the current project. "
+        "If no team name is provided, the default team will be shown."
+    )
 
     team_parser = auth_subparser.add_parser(
         "team",
@@ -102,13 +107,15 @@ def add_auth_parser(subparsers, parent_parser) -> None:
     team_parser.add_argument(
         "team_name",
         action="store",
+        default=None,
+        nargs="?",
         help="Name of the team to set as default",
     )
     team_parser.add_argument(
-        "--global",
+        "--local",
         action="store_true",
         default=False,
-        help="Set team globally for all projects",
+        help="Set team locally for the current project",
     )
 
     auth_token_help = "View Studio authentication token"  # noqa: S105

datachain/cli/parser/utils.py

@@ -1,5 +1,4 @@
 from argparse import Action, ArgumentParser, ArgumentTypeError, HelpFormatter
-from typing import Union
 
 from datachain.cli.utils import CommaSeparatedArgs
 
@@ -44,7 +43,7 @@ def parse_find_column(column: str) -> str:
     )
 
 
-def add_sources_arg(parser: ArgumentParser, nargs: Union[str, int] = "+") -> Action:
+def add_sources_arg(parser: ArgumentParser, nargs: str | int = "+") -> Action:
     return parser.add_argument(
         "sources",
         type=str,

datachain/cli/utils.py

@@ -1,6 +1,5 @@
 import logging
-from argparse import SUPPRESS, Action, ArgumentError, Namespace, _AppendAction
-from typing import Optional
+from argparse import SUPPRESS, Action, Namespace, _AppendAction
 
 from datachain.error import DataChainError
 
@@ -64,18 +63,6 @@ class CommaSeparatedArgs(_AppendAction):  # pylint: disable=protected-access
         setattr(namespace, self.dest, list(dict.fromkeys(items)))
 
 
-class KeyValueArgs(_AppendAction):  # pylint: disable=protected-access
-    def __call__(self, parser, namespace, values, option_string=None):
-        items = getattr(namespace, self.dest) or {}
-        for raw_value in filter(bool, values):
-            key, sep, value = raw_value.partition("=")
-            if not key or not sep or value == "":
-                raise ArgumentError(self, f"expected 'key=value', got {raw_value!r}")
-            items[key.strip()] = value
-
-        setattr(namespace, self.dest, items)
-
-
 def get_logging_level(args: Namespace) -> int:
     if args.quiet:
         return logging.CRITICAL
@@ -84,7 +71,7 @@ def get_logging_level(args: Namespace) -> int:
     return logging.INFO
 
 
-def determine_flavors(studio: bool, local: bool, all: bool, token: Optional[str]):
+def determine_flavors(studio: bool, local: bool, all: bool, token: str | None):
     if studio and not token:
         raise DataChainError(
             "Not logged in to Studio. Log in with 'datachain auth login'."

datachain/client/azure.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 from urllib.parse import parse_qs, urlsplit, urlunsplit
 
 from adlfs import AzureBlobFileSystem
@@ -15,7 +15,7 @@ class AzureClient(Client):
     protocol = "az"
 
     def info_to_file(self, v: dict[str, Any], path: str) -> File:
-        version_id = v.get("version_id")
+        version_id = v.get("version_id") if self._is_version_aware() else None
         return File(
             source=self.uri,
             path=path,
@@ -65,7 +65,7 @@ class AzureClient(Client):
                         if entries:
                             await result_queue.put(entries)
                             pbar.update(len(entries))
-                    if not found:
+                    if not found and prefix:
                         raise FileNotFoundError(
                             f"Unable to resolve remote path: {prefix}"
                         )
@@ -73,7 +73,7 @@ class AzureClient(Client):
             result_queue.put_nowait(None)
 
     @classmethod
-    def version_path(cls, path: str, version_id: Optional[str]) -> str:
+    def version_path(cls, path: str, version_id: str | None) -> str:
         parts = list(urlsplit(path))
         query = parse_qs(parts[3])
         if "versionid" in query:

datachain/client/fsspec.py

@@ -10,15 +10,7 @@ from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Iterator, Sequence
 from datetime import datetime
 from shutil import copy2
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    BinaryIO,
-    ClassVar,
-    NamedTuple,
-    Optional,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar, NamedTuple
 from urllib.parse import urlparse
 
 from dvc_objects.fs.system import reflink
@@ -44,11 +36,12 @@ FETCH_WORKERS = 100
 DELIMITER = "/"  # Path delimiter.
 
 DATA_SOURCE_URI_PATTERN = re.compile(r"^[\w]+:\/\/.*$")
+CLOUD_STORAGE_PROTOCOLS = {"s3", "gs", "az", "hf"}
 
-ResultQueue = asyncio.Queue[Optional[Sequence["File"]]]
+ResultQueue = asyncio.Queue[Sequence["File"] | None]
 
 
-def _is_win_local_path(uri: str) -> bool:
+def is_win_local_path(uri: str) -> bool:
     if sys.platform == "win32":
         if len(uri) >= 1 and uri[0] == "\\":
             return True
@@ -62,10 +55,20 @@ def _is_win_local_path(uri: str) -> bool:
     return False
 
 
+def is_cloud_uri(uri: str) -> bool:
+    protocol = urlparse(uri).scheme
+    return protocol in CLOUD_STORAGE_PROTOCOLS
+
+
+def get_cloud_schemes() -> list[str]:
+    """Get list of cloud storage scheme prefixes."""
+    return [f"{p}://" for p in CLOUD_STORAGE_PROTOCOLS]
+
+
 class Bucket(NamedTuple):
     name: str
     uri: "StorageURI"
-    created: Optional[datetime]
+    created: datetime | None
 
 
 class Client(ABC):
@@ -77,21 +80,22 @@ class Client(ABC):
     def __init__(self, name: str, fs_kwargs: dict[str, Any], cache: Cache) -> None:
         self.name = name
         self.fs_kwargs = fs_kwargs
-        self._fs: Optional[AbstractFileSystem] = None
+        self._fs: AbstractFileSystem | None = None
         self.cache = cache
         self.uri = self.get_uri(self.name)
 
     @staticmethod
-    def get_implementation(url: Union[str, os.PathLike[str]]) -> type["Client"]:
+    def get_implementation(url: str | os.PathLike[str]) -> type["Client"]:  # noqa: PLR0911
         from .azure import AzureClient
         from .gcs import GCSClient
         from .hf import HfClient
+        from .http import HTTPClient, HTTPSClient
         from .local import FileClient
         from .s3 import ClientS3
 
         protocol = urlparse(os.fspath(url)).scheme
 
-        if not protocol or _is_win_local_path(os.fspath(url)):
+        if not protocol or is_win_local_path(os.fspath(url)):
             return FileClient
         if protocol == ClientS3.protocol:
             return ClientS3
@@ -103,9 +107,18 @@ class Client(ABC):
             return FileClient
         if protocol == HfClient.protocol:
             return HfClient
+        if protocol == HTTPClient.protocol:
+            return HTTPClient
+        if protocol == HTTPSClient.protocol:
+            return HTTPSClient
 
         raise NotImplementedError(f"Unsupported protocol: {protocol}")
 
+    @classmethod
+    def path_to_uri(cls, path: str) -> str:
+        """Convert a path-like object to a URI. Default: identity."""
+        return path
+
     @staticmethod
     def is_data_source_uri(name: str) -> bool:
         # Returns True if name is one of supported data sources URIs, e.g s3 bucket
@@ -118,9 +131,7 @@ class Client(ABC):
         return cls.get_uri(storage_name), rel_path
 
     @staticmethod
-    def get_client(
-        source: Union[str, os.PathLike[str]], cache: Cache, **kwargs
-    ) -> "Client":
+    def get_client(source: str | os.PathLike[str], cache: Cache, **kwargs) -> "Client":
         cls = Client.get_implementation(source)
         storage_url, _ = cls.split_url(os.fspath(source))
         if os.name == "nt":
@@ -136,7 +147,7 @@ class Client(ABC):
         return fs
 
     @classmethod
-    def version_path(cls, path: str, version_id: Optional[str]) -> str:
+    def version_path(cls, path: str, version_id: str | None) -> str:
         return path
 
     @classmethod
@@ -207,24 +218,25 @@ class Client(ABC):
         )
 
     async def get_current_etag(self, file: "File") -> str:
+        file_path = file.get_path_normalized()
         kwargs = {}
-        if getattr(self.fs, "version_aware", False):
+        if self._is_version_aware():
             kwargs["version_id"] = file.version
         info = await self.fs._info(
-            self.get_full_path(file.path, file.version), **kwargs
+            self.get_full_path(file_path, file.version), **kwargs
         )
-        return self.info_to_file(info, file.path).etag
+        return self.info_to_file(info, file_path).etag
 
-    def get_file_info(self, path: str, version_id: Optional[str] = None) -> "File":
+    def get_file_info(self, path: str, version_id: str | None = None) -> "File":
         info = self.fs.info(self.get_full_path(path, version_id), version_id=version_id)
         return self.info_to_file(info, path)
 
-    async def get_size(self, path: str, version_id: Optional[str] = None) -> int:
+    async def get_size(self, path: str, version_id: str | None = None) -> int:
         return await self.fs._size(
             self.version_path(path, version_id), version_id=version_id
         )
 
-    async def get_file(self, lpath, rpath, callback, version_id: Optional[str] = None):
+    async def get_file(self, lpath, rpath, callback, version_id: str | None = None):
         return await self.fs._get_file(
             self.version_path(lpath, version_id),
             rpath,
@@ -326,15 +338,19 @@ class Client(ABC):
         """
         return not (key.startswith("/") or key.endswith("/") or "//" in key)
 
+    def _is_version_aware(self) -> bool:
+        return getattr(self.fs, "version_aware", False)
+
     async def ls_dir(self, path):
-        if getattr(self.fs, "version_aware", False):
+        kwargs = {}
+        if self._is_version_aware():
             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]
 
-    def get_full_path(self, rel_path: str, version_id: Optional[str] = None) -> str:
+    def get_full_path(self, rel_path: str, version_id: str | None = None) -> str:
         return self.version_path(f"{self.PREFIX}{self.name}/{rel_path}", version_id)
 
     @abstractmethod
@@ -382,7 +398,8 @@ class Client(ABC):
             return open(cache_path, mode="rb")
         assert not file.location
         return FileWrapper(
-            self.fs.open(self.get_full_path(file.path, file.version)), cb
+            self.fs.open(self.get_full_path(file.get_path_normalized(), file.version)),
+            cb,
         )  # type: ignore[return-value]
 
     def upload(self, data: bytes, path: str) -> "File":