datachain 0.8.2__py3-none-any.whl → 0.8.4__py3-none-any.whl

datachain/cli/commands/datasets.py

@@ -0,0 +1,129 @@
+import sys
+from typing import TYPE_CHECKING, Optional
+
+from tabulate import tabulate
+
+from datachain import utils
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+from datachain.cli.utils import determine_flavors
+from datachain.config import Config
+from datachain.error import DatasetNotFoundError
+
+
+def list_datasets(
+    catalog: "Catalog",
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
+):
+    from datachain.studio import list_datasets
+
+    token = Config().read().get("studio", {}).get("token")
+    all, local, studio = determine_flavors(studio, local, all, token)
+
+    local_datasets = set(list_datasets_local(catalog)) if all or local else set()
+    studio_datasets = (
+        set(list_datasets(team=team)) if (all or studio) and token else set()
+    )
+
+    rows = [
+        _datasets_tabulate_row(
+            name=name,
+            version=version,
+            both=(all or (local and studio)) and token,
+            local=(name, version) in local_datasets,
+            studio=(name, version) in studio_datasets,
+        )
+        for name, version in local_datasets.union(studio_datasets)
+    ]
+
+    print(tabulate(rows, headers="keys"))
+
+
+def list_datasets_local(catalog: "Catalog"):
+    for d in catalog.ls_datasets():
+        for v in d.versions:
+            yield (d.name, v.version)
+
+
+def _datasets_tabulate_row(name, version, both, local, studio):
+    row = {
+        "Name": name,
+        "Version": version,
+    }
+    if both:
+        row["Studio"] = "\u2714" if studio else "\u2716"
+        row["Local"] = "\u2714" if local else "\u2716"
+    return row
+
+
+def rm_dataset(
+    catalog: "Catalog",
+    name: str,
+    version: Optional[int] = None,
+    force: Optional[bool] = False,
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
+):
+    from datachain.studio import remove_studio_dataset
+
+    token = Config().read().get("studio", {}).get("token")
+    all, local, studio = determine_flavors(studio, local, all, token)
+
+    if all or local:
+        try:
+            catalog.remove_dataset(name, version=version, force=force)
+        except DatasetNotFoundError:
+            print("Dataset not found in local", file=sys.stderr)
+
+    if (all or studio) and token:
+        remove_studio_dataset(team, name, version, force)
+
+
+def edit_dataset(
+    catalog: "Catalog",
+    name: str,
+    new_name: Optional[str] = None,
+    description: Optional[str] = None,
+    labels: Optional[list[str]] = None,
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
+):
+    from datachain.studio import edit_studio_dataset
+
+    token = Config().read().get("studio", {}).get("token")
+    all, local, studio = determine_flavors(studio, local, all, token)
+
+    if all or local:
+        try:
+            catalog.edit_dataset(name, new_name, description, labels)
+        except DatasetNotFoundError:
+            print("Dataset not found in local", file=sys.stderr)
+
+    if (all or studio) and token:
+        edit_studio_dataset(team, name, new_name, description, labels)
+
+
+def dataset_stats(
+    catalog: "Catalog",
+    name: str,
+    version: int,
+    show_bytes=False,
+    si=False,
+):
+    stats = catalog.dataset_stats(name, version)
+
+    if stats:
+        print(f"Number of objects: {stats.num_objects}")
+        if show_bytes:
+            print(f"Total objects size: {stats.size}")
+        else:
+            print(f"Total objects size: {utils.sizeof_fmt(stats.size, si=si): >7}")

datachain/cli/commands/du.py

@@ -0,0 +1,14 @@
+from typing import TYPE_CHECKING
+
+from datachain import utils
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+
+def du(catalog: "Catalog", sources, show_bytes=False, si=False, **kwargs):
+    for path, size in catalog.du(sources, **kwargs):
+        if show_bytes:
+            print(f"{size} {path}")
+        else:
+            print(f"{utils.sizeof_fmt(size, si=si): >7} {path}")

datachain/cli/commands/index.py

@@ -0,0 +1,12 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+
+def index(
+    catalog: "Catalog",
+    sources,
+    **kwargs,
+):
+    catalog.index(sources, **kwargs)

datachain/cli/commands/ls.py

@@ -0,0 +1,169 @@
+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 datachain.cli.utils import determine_flavors
+from datachain.config import Config
+
+
+def ls(
+    sources,
+    long: bool = False,
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
+    **kwargs,
+):
+    token = Config().read().get("studio", {}).get("token")
+    all, local, studio = determine_flavors(studio, local, all, token)
+
+    if all or local:
+        ls_local(sources, long=long, **kwargs)
+
+    if (all or studio) and token:
+        ls_remote(sources, long=long, team=team)
+
+
+def ls_local(
+    sources,
+    long: bool = False,
+    catalog: Optional["Catalog"] = None,
+    client_config=None,
+    **kwargs,
+):
+    from datachain import DataChain
+
+    if catalog is None:
+        from datachain.catalog import get_catalog
+
+        catalog = get_catalog(client_config=client_config)
+    if sources:
+        actual_sources = list(ls_urls(sources, catalog=catalog, long=long, **kwargs))
+        if len(actual_sources) == 1:
+            for _, entries in actual_sources:
+                for entry in entries:
+                    print(format_ls_entry(entry))
+        else:
+            first = True
+            for source, entries in actual_sources:
+                # print a newline between directory listings
+                if first:
+                    first = False
+                else:
+                    print()
+                if source:
+                    print(f"{source}:")
+                for entry in entries:
+                    print(format_ls_entry(entry))
+    else:
+        chain = DataChain.listings()
+        for ls in chain.collect("listing"):
+            print(format_ls_entry(f"{ls.uri}@v{ls.version}"))  # type: ignore[union-attr]
+
+
+def format_ls_entry(entry: str) -> str:
+    if entry.endswith("/") or not entry:
+        entry = shlex.quote(entry[:-1])
+        return f"{entry}/"
+    return shlex.quote(entry)
+
+
+def ls_remote(
+    paths: Iterable[str],
+    long: bool = False,
+    team: Optional[str] = None,
+):
+    from datachain.node import long_line_str
+    from datachain.remote.studio import StudioClient
+
+    client = StudioClient(team=team)
+    first = True
+    for path, response in client.ls(paths):
+        if not first:
+            print()
+        if not response.ok or response.data is None:
+            print(f"{path}:\n  Error: {response.message}\n")
+            continue
+
+        print(f"{path}:")
+        if long:
+            for row in response.data:
+                entry = long_line_str(
+                    row["name"] + ("/" if row["dir_type"] else ""),
+                    row["last_modified"],
+                )
+                print(format_ls_entry(entry))
+        else:
+            for row in response.data:
+                entry = row["name"] + ("/" if row["dir_type"] else "")
+                print(format_ls_entry(entry))
+        first = False
+
+
+def ls_urls(
+    sources,
+    catalog: "Catalog",
+    long: bool = False,
+    **kwargs,
+) -> Iterator[tuple[str, Iterator[str]]]:
+    curr_dir = None
+    value_iterables = []
+    for next_dir, values in _ls_urls_flat(sources, long, catalog, **kwargs):
+        if curr_dir is None or next_dir == curr_dir:  # type: ignore[unreachable]
+            value_iterables.append(values)
+        else:
+            yield curr_dir, chain(*value_iterables)  # type: ignore[unreachable]
+            value_iterables = [values]
+        curr_dir = next_dir
+    if curr_dir is not None:
+        yield curr_dir, chain(*value_iterables)
+
+
+def _node_data_to_ls_values(row, long_format=False):
+    from datachain.node import DirType, long_line_str
+
+    name = row[0]
+    is_dir = row[1] == DirType.DIR
+    ending = "/" if is_dir else ""
+    value = name + ending
+    if long_format:
+        last_modified = row[2]
+        timestamp = last_modified if not is_dir else None
+        return long_line_str(value, timestamp)
+    return value
+
+
+def _ls_urls_flat(
+    sources,
+    long: bool,
+    catalog: "Catalog",
+    **kwargs,
+) -> Iterator[tuple[str, Iterator[str]]]:
+    from datachain.client import Client
+    from datachain.node import long_line_str
+
+    for source in sources:
+        client_cls = Client.get_implementation(source)
+        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)
+            else:
+                values = (b.name for b in buckets)
+            yield source, values
+        else:
+            found = False
+            fields = ["name", "dir_type"]
+            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)
+                found = True
+                yield data_source.dirname(), values
+            if not found:
+                raise FileNotFoundError(f"No such file or directory: {source}")

datachain/cli/commands/misc.py

@@ -0,0 +1,28 @@
+from typing import TYPE_CHECKING
+
+import shtab
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+
+def clear_cache(catalog: "Catalog"):
+    catalog.cache.clear()
+
+
+def garbage_collect(catalog: "Catalog"):
+    temp_tables = catalog.get_temp_table_names()
+    if not temp_tables:
+        print("Nothing to clean up.")
+    else:
+        print(f"Garbage collecting {len(temp_tables)} tables.")
+        catalog.cleanup_tables(temp_tables)
+
+
+def completion(shell: str) -> str:
+    from datachain.cli import get_parser
+
+    return shtab.complete(
+        get_parser(),
+        shell=shell,
+    )

datachain/cli/commands/query.py

@@ -0,0 +1,53 @@
+import os
+import sys
+import traceback
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+
+def query(
+    catalog: "Catalog",
+    script: str,
+    parallel: Optional[int] = None,
+    params: Optional[dict[str, str]] = None,
+) -> None:
+    from datachain.data_storage import JobQueryType, JobStatus
+
+    with open(script, encoding="utf-8") as f:
+        script_content = f.read()
+
+    if parallel is not None:
+        # This also sets this environment variable for any subprocesses
+        os.environ["DATACHAIN_SETTINGS_PARALLEL"] = str(parallel)
+
+    python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+    python_executable = sys.executable
+
+    job_id = catalog.metastore.create_job(
+        name=os.path.basename(script),
+        query=script_content,
+        query_type=JobQueryType.PYTHON,
+        python_version=python_version,
+        params=params,
+    )
+
+    try:
+        catalog.query(
+            script_content,
+            python_executable=python_executable,
+            params=params,
+            job_id=job_id,
+        )
+    except Exception as e:
+        error_message = str(e)
+        error_stack = traceback.format_exc()
+        catalog.metastore.set_job_status(
+            job_id,
+            JobStatus.FAILED,
+            error_message=error_message,
+            error_stack=error_stack,
+        )
+        raise
+    catalog.metastore.set_job_status(job_id, JobStatus.COMPLETE)

datachain/cli/commands/show.py

@@ -0,0 +1,38 @@
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+
+def show(
+    catalog: "Catalog",
+    name: str,
+    version: Optional[int] = None,
+    limit: int = 10,
+    offset: int = 0,
+    columns: Sequence[str] = (),
+    no_collapse: bool = False,
+    schema: bool = False,
+) -> None:
+    from datachain import Session
+    from datachain.lib.dc import DataChain
+    from datachain.query.dataset import DatasetQuery
+    from datachain.utils import show_records
+
+    dataset = catalog.get_dataset(name)
+    dataset_version = dataset.get_version(version or dataset.latest_version)
+
+    query = (
+        DatasetQuery(name=name, version=version, catalog=catalog)
+        .select(*columns)
+        .limit(limit)
+        .offset(offset)
+    )
+    records = query.to_db_records()
+    show_records(records, collapse_columns=not no_collapse)
+    if schema and dataset_version.feature_schema:
+        print("\nSchema:")
+        session = Session.get(catalog=catalog)
+        dc = DataChain.from_dataset(name=name, version=version, session=session)
+        dc.print_schema()