datachain 0.6.7__tar.gz → 0.6.8__tar.gz

{datachain-0.6.7 → datachain-0.6.8}/.github/workflows/tests-studio.yml

@@ -32,7 +32,7 @@ jobs:
           POSTGRES_DB: database
           POSTGRES_HOST_AUTH_METHOD: trust
       clickhouse:
-        image: clickhouse/clickhouse-server:24
+        image: clickhouse/clickhouse-server:24.6
         ports:
           - 8123:8123
           - 9010:9000

{datachain-0.6.7/src/datachain.egg-info → datachain-0.6.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.6.7
+Version: 0.6.8
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -45,6 +45,7 @@ Requires-Dist: huggingface_hub
 Requires-Dist: iterative-telemetry>=0.0.9
 Requires-Dist: platformdirs
 Requires-Dist: dvc-studio-client<1,>=0.21
+Requires-Dist: tabulate
 Provides-Extra: docs
 Requires-Dist: mkdocs>=1.5.2; extra == "docs"
 Requires-Dist: mkdocs-gen-files>=0.5.0; extra == "docs"
@@ -87,6 +88,7 @@ Requires-Dist: types-python-dateutil; extra == "dev"
 Requires-Dist: types-pytz; extra == "dev"
 Requires-Dist: types-PyYAML; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
+Requires-Dist: types-tabulate; extra == "dev"
 Provides-Extra: examples
 Requires-Dist: datachain[tests]; extra == "examples"
 Requires-Dist: numpy<2,>=1; extra == "examples"

{datachain-0.6.7 → datachain-0.6.8}/pyproject.toml

@@ -47,7 +47,8 @@ dependencies = [
   "huggingface_hub",
   "iterative-telemetry>=0.0.9",
   "platformdirs",
-  "dvc-studio-client>=0.21,<1"
+  "dvc-studio-client>=0.21,<1",
+  "tabulate"
 ]
 
 [project.optional-dependencies]
@@ -98,7 +99,8 @@ dev = [
   "types-python-dateutil",
   "types-pytz",
   "types-PyYAML",
-  "types-requests"
+  "types-requests",
+  "types-tabulate"
 ]
 examples = [
   "datachain[tests]",

{datachain-0.6.7 → datachain-0.6.8}/src/datachain/__init__.py

@@ -1,4 +1,4 @@
-from datachain.lib import func
+from datachain.lib import func, models
 from datachain.lib.data_model import DataModel, DataType, is_chain_type
 from datachain.lib.dc import C, Column, DataChain, Sys
 from datachain.lib.file import (
@@ -38,5 +38,6 @@ __all__ = [
     "func",
     "is_chain_type",
     "metrics",
+    "models",
     "param",
 ]

{datachain-0.6.7 → datachain-0.6.8}/src/datachain/cli.py

@@ -4,18 +4,21 @@ import shlex
 import sys
 import traceback
 from argparse import Action, ArgumentParser, ArgumentTypeError, Namespace
-from collections.abc import Iterable, Iterator, Mapping, Sequence
+from collections.abc import Iterable, Iterator, Sequence
 from importlib.metadata import PackageNotFoundError, version
 from itertools import chain
 from multiprocessing import freeze_support
 from typing import TYPE_CHECKING, Optional, Union
 
 import shtab
+from tabulate import tabulate
 
 from datachain import Session, utils
 from datachain.cli_utils import BooleanOptionalAction, CommaSeparatedArgs, KeyValueArgs
+from datachain.config import Config
+from datachain.error import DataChainError
 from datachain.lib.dc import DataChain
-from datachain.studio import process_studio_cli_args
+from datachain.studio import list_datasets, process_studio_cli_args
 from datachain.telemetry import telemetry
 
 if TYPE_CHECKING:
@@ -416,7 +419,36 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="Dataset labels",
     )
 
-    subp.add_parser("ls-datasets", parents=[parent_parser], description="List datasets")
+    datasets_parser = subp.add_parser(
+        "datasets", parents=[parent_parser], description="List datasets"
+    )
+    datasets_parser.add_argument(
+        "--studio",
+        action="store_true",
+        default=False,
+        help="List the files in the Studio",
+    )
+    datasets_parser.add_argument(
+        "-L",
+        "--local",
+        action="store_true",
+        default=False,
+        help="List local files only",
+    )
+    datasets_parser.add_argument(
+        "-a",
+        "--all",
+        action="store_true",
+        default=True,
+        help="List all files including hidden files",
+    )
+    datasets_parser.add_argument(
+        "--team",
+        action="store",
+        default=None,
+        help="The team to list datasets for. By default, it will use team from config.",
+    )
+
     rm_dataset_parser = subp.add_parser(
         "rm-dataset", parents=[parent_parser], description="Removes dataset"
     )
@@ -474,10 +506,30 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="List files in the long format",
     )
     parse_ls.add_argument(
-        "--remote",
+        "--studio",
+        action="store_true",
+        default=False,
+        help="List the files in the Studio",
+    )
+    parse_ls.add_argument(
+        "-L",
+        "--local",
+        action="store_true",
+        default=False,
+        help="List local files only",
+    )
+    parse_ls.add_argument(
+        "-a",
+        "--all",
+        action="store_true",
+        default=True,
+        help="List all files including hidden files",
+    )
+    parse_ls.add_argument(
+        "--team",
         action="store",
-        default="",
-        help="Name of remote to use",
+        default=None,
+        help="The team to list datasets for. By default, it will use team from config.",
     )
 
     parse_du = subp.add_parser(
@@ -758,11 +810,12 @@ def format_ls_entry(entry: str) -> str:
 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()
+    client = StudioClient(team=team)
     first = True
     for path, response in client.ls(paths):
         if not first:
@@ -789,28 +842,66 @@ def ls_remote(
 def ls(
     sources,
     long: bool = False,
-    remote: str = "",
-    config: Optional[Mapping[str, str]] = None,
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
     **kwargs,
 ):
-    if config is None:
-        from .config import Config
+    token = Config().read().get("studio", {}).get("token")
+    all, local, studio = _determine_flavors(studio, local, all, token)
 
-        config = Config().get_remote_config(remote=remote)
-    remote_type = config["type"]
-    if remote_type == "local":
+    if all or local:
         ls_local(sources, long=long, **kwargs)
-    else:
-        ls_remote(
-            sources,
-            long=long,
+
+    if (all or studio) and token:
+        ls_remote(sources, long=long, team=team)
+
+
+def datasets(
+    catalog: "Catalog",
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
+):
+    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 ls_datasets(catalog: "Catalog"):
+def list_datasets_local(catalog: "Catalog"):
     for d in catalog.ls_datasets():
         for v in d.versions:
-            print(f"{d.name} (v{v.version})")
+            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(
@@ -953,6 +1044,20 @@ def completion(shell: str) -> str:
     )
 
 
+def _determine_flavors(studio: bool, local: bool, all: bool, token: Optional[str]):
+    if studio and not token:
+        raise DataChainError(
+            "Not logged in to Studio. Log in with 'datachain studio login'."
+        )
+
+    if local or studio:
+        all = False
+
+    all = all and not (local or studio)
+
+    return all, local, studio
+
+
 def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0912, PLR0915
     # Required for Windows multiprocessing support
     freeze_support()
@@ -1032,12 +1137,21 @@ def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0912, PLR09
             ls(
                 args.sources,
                 long=bool(args.long),
-                remote=args.remote,
+                studio=args.studio,
+                local=args.local,
+                all=args.all,
+                team=args.team,
                 update=bool(args.update),
                 client_config=client_config,
             )
-        elif args.command == "ls-datasets":
-            ls_datasets(catalog)
+        elif args.command == "datasets":
+            datasets(
+                catalog=catalog,
+                studio=args.studio,
+                local=args.local,
+                all=args.all,
+                team=args.team,
+            )
         elif args.command == "show":
             show(
                 catalog,

{datachain-0.6.7 → datachain-0.6.8}/src/datachain/lib/file.py

@@ -20,9 +20,6 @@ from PIL import Image
 from pyarrow.dataset import dataset
 from pydantic import Field, field_validator
 
-if TYPE_CHECKING:
-    from typing_extensions import Self
-
 from datachain.client.fileslice import FileSlice
 from datachain.lib.data_model import DataModel
 from datachain.lib.utils import DataChainError

datachain-0.6.8/src/datachain/lib/models/__init__.py

@@ -0,0 +1,5 @@
+from . import yolo
+from .bbox import BBox
+from .pose import Pose, Pose3D
+
+__all__ = ["BBox", "Pose", "Pose3D", "yolo"]

datachain-0.6.8/src/datachain/lib/models/bbox.py

@@ -0,0 +1,45 @@
+from typing import Optional
+
+from pydantic import Field
+
+from datachain.lib.data_model import DataModel
+
+
+class BBox(DataModel):
+    """
+    A data model for representing bounding boxes.
+
+    Attributes:
+        title (str): The title of the bounding box.
+        x1 (float): The x-coordinate of the top-left corner of the bounding box.
+        y1 (float): The y-coordinate of the top-left corner of the bounding box.
+        x2 (float): The x-coordinate of the bottom-right corner of the bounding box.
+        y2 (float): The y-coordinate of the bottom-right corner of the bounding box.
+
+    The bounding box is defined by two points:
+        - (x1, y1): The top-left corner of the box.
+        - (x2, y2): The bottom-right corner of the box.
+    """
+
+    title: str = Field(default="")
+    x1: float = Field(default=0)
+    y1: float = Field(default=0)
+    x2: float = Field(default=0)
+    y2: float = Field(default=0)
+
+    @staticmethod
+    def from_xywh(bbox: list[float], title: Optional[str] = None) -> "BBox":
+        """
+        Converts a bounding box in (x, y, width, height) format
+        to a BBox data model instance.
+
+        Args:
+            bbox (list[float]): A bounding box, represented as a list
+                                of four floats [x, y, width, height].
+
+        Returns:
+            BBox2D: An instance of the BBox data model.
+        """
+        assert len(bbox) == 4, f"Bounding box must have 4 elements, got f{len(bbox)}"
+        x, y, w, h = bbox
+        return BBox(title=title or "", x1=x, y1=y, x2=x + w, y2=y + h)

datachain-0.6.8/src/datachain/lib/models/pose.py

@@ -0,0 +1,37 @@
+from pydantic import Field
+
+from datachain.lib.data_model import DataModel
+
+
+class Pose(DataModel):
+    """
+    A data model for representing pose keypoints.
+
+    Attributes:
+        x (list[float]): The x-coordinates of the keypoints.
+        y (list[float]): The y-coordinates of the keypoints.
+
+    The keypoints are represented as lists of x and y coordinates, where each index
+    corresponds to a specific body part.
+    """
+
+    x: list[float] = Field(default=None)
+    y: list[float] = Field(default=None)
+
+
+class Pose3D(DataModel):
+    """
+    A data model for representing 3D pose keypoints.
+
+    Attributes:
+        x (list[float]): The x-coordinates of the keypoints.
+        y (list[float]): The y-coordinates of the keypoints.
+        visible (list[float]): The visibility of the keypoints.
+
+    The keypoints are represented as lists of x, y, and visibility values,
+    where each index corresponds to a specific body part.
+    """
+
+    x: list[float] = Field(default=None)
+    y: list[float] = Field(default=None)
+    visible: list[float] = Field(default=None)

datachain-0.6.8/src/datachain/lib/models/yolo.py

@@ -0,0 +1,39 @@
+"""
+This module contains the YOLO models.
+
+YOLO stands for "You Only Look Once", a family of object detection models that
+are designed to be fast and accurate. The models are trained to detect objects
+in images by dividing the image into a grid and predicting the bounding boxes
+and class probabilities for each grid cell.
+
+More information about YOLO can be found here:
+- https://pjreddie.com/darknet/yolo/
+- https://docs.ultralytics.com/
+"""
+
+
+class PoseBodyPart:
+    """
+    An enumeration of body parts for YOLO pose keypoints.
+
+    More information about the body parts can be found here:
+    https://docs.ultralytics.com/tasks/pose/
+    """
+
+    nose = 0
+    left_eye = 1
+    right_eye = 2
+    left_ear = 3
+    right_ear = 4
+    left_shoulder = 5
+    right_shoulder = 6
+    left_elbow = 7
+    right_elbow = 8
+    left_wrist = 9
+    right_wrist = 10
+    left_hip = 11
+    right_hip = 12
+    left_knee = 13
+    right_knee = 14
+    left_ankle = 15
+    right_ankle = 16

{datachain-0.6.7 → datachain-0.6.8}/src/datachain/remote/studio.py

@@ -131,6 +131,12 @@ class StudioClient:
             timeout=self.timeout,
         )
         ok = response.ok
+        if not ok:
+            if response.status_code == 403:
+                message = f"Not authorized for the team {self.team}"
+                raise DataChainError(message)
+            logger.error("Got bad response from Studio")
+
         content = msgpack.unpackb(response.content, ext_hook=self._unpacker_hook)
         response_data = content.get("data")
         if ok and response_data is None:
@@ -177,8 +183,12 @@ class StudioClient:
                 response.content.decode("utf-8"),
             )
             if response.status_code == 403:
-                message = "Not authorized"
+                message = f"Not authorized for the team {self.team}"
             else:
+                logger.error(
+                    "Got bad response from Studio, content is %s",
+                    response.content.decode("utf-8"),
+                )
                 message = data.get("message", "")
         else:
             message = ""
@@ -214,7 +224,7 @@ class StudioClient:
         # to handle cases where a path will be expanded (i.e. globs)
         response: Response[LsData]
         for path in paths:
-            response = self._send_request_msgpack("ls", {"source": path})
+            response = self._send_request_msgpack("datachain/ls", {"source": path})
             yield path, response
 
     def ls_datasets(self) -> Response[LsData]:

{datachain-0.6.7 → datachain-0.6.8}/src/datachain/studio.py

@@ -1,8 +1,11 @@
 import os
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
+
+from tabulate import tabulate
 
 from datachain.catalog.catalog import raise_remote_error
 from datachain.config import Config, ConfigLevel
+from datachain.dataset import QUERY_DATASET_PREFIX
 from datachain.error import DataChainError
 from datachain.remote.studio import StudioClient
 from datachain.utils import STUDIO_URL
@@ -24,7 +27,13 @@ def process_studio_cli_args(args: "Namespace"):
     if args.cmd == "token":
         return token()
     if args.cmd == "datasets":
-        return list_datasets(args)
+        rows = [
+            {"Name": name, "Version": version}
+            for name, version in list_datasets(args.team)
+        ]
+        print(tabulate(rows, headers="keys"))
+        return 0
+
     if args.cmd == "team":
         return set_team(args)
     raise DataChainError(f"Unknown command '{args.cmd}'.")
@@ -103,19 +112,22 @@ def token():
     print(token)
 
 
-def list_datasets(args: "Namespace"):
-    client = StudioClient(team=args.team)
+def list_datasets(team: Optional[str] = None):
+    client = StudioClient(team=team)
     response = client.ls_datasets()
     if not response.ok:
         raise_remote_error(response.message)
     if not response.data:
-        print("No datasets found.")
         return
+
     for d in response.data:
         name = d.get("name")
+        if name and name.startswith(QUERY_DATASET_PREFIX):
+            continue
+
         for v in d.get("versions", []):
             version = v.get("version")
-            print(f"{name} (v{version})")
+            yield (name, version)
 
 
 def save_config(hostname, token):

{datachain-0.6.7 → datachain-0.6.8/src/datachain.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.6.7
+Version: 0.6.8
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -45,6 +45,7 @@ Requires-Dist: huggingface_hub
 Requires-Dist: iterative-telemetry>=0.0.9
 Requires-Dist: platformdirs
 Requires-Dist: dvc-studio-client<1,>=0.21
+Requires-Dist: tabulate
 Provides-Extra: docs
 Requires-Dist: mkdocs>=1.5.2; extra == "docs"
 Requires-Dist: mkdocs-gen-files>=0.5.0; extra == "docs"
@@ -87,6 +88,7 @@ Requires-Dist: types-python-dateutil; extra == "dev"
 Requires-Dist: types-pytz; extra == "dev"
 Requires-Dist: types-PyYAML; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
+Requires-Dist: types-tabulate; extra == "dev"
 Provides-Extra: examples
 Requires-Dist: datachain[tests]; extra == "examples"
 Requires-Dist: numpy<2,>=1; extra == "examples"

{datachain-0.6.7 → datachain-0.6.8}/src/datachain.egg-info/SOURCES.txt

@@ -129,6 +129,10 @@ src/datachain/lib/convert/values_to_tuples.py
 src/datachain/lib/func/__init__.py
 src/datachain/lib/func/aggregate.py
 src/datachain/lib/func/func.py
+src/datachain/lib/models/__init__.py
+src/datachain/lib/models/bbox.py
+src/datachain/lib/models/pose.py
+src/datachain/lib/models/yolo.py
 src/datachain/query/__init__.py
 src/datachain/query/batch.py
 src/datachain/query/dataset.py
@@ -237,6 +241,7 @@ tests/unit/lib/test_file.py
 tests/unit/lib/test_hf.py
 tests/unit/lib/test_image.py
 tests/unit/lib/test_listing_info.py
+tests/unit/lib/test_models.py
 tests/unit/lib/test_schema.py
 tests/unit/lib/test_signal_schema.py
 tests/unit/lib/test_sql_to_python.py

{datachain-0.6.7 → datachain-0.6.8}/src/datachain.egg-info/requires.txt

@@ -27,6 +27,7 @@ huggingface_hub
 iterative-telemetry>=0.0.9
 platformdirs
 dvc-studio-client<1,>=0.21
+tabulate
 
 [dev]
 datachain[docs,tests]
@@ -35,6 +36,7 @@ types-python-dateutil
 types-pytz
 types-PyYAML
 types-requests
+types-tabulate
 
 [docs]
 mkdocs>=1.5.2

{datachain-0.6.7 → datachain-0.6.8}/tests/conftest.py

@@ -14,6 +14,7 @@ from upath.implementations.cloud import CloudPath
 from datachain.catalog import Catalog
 from datachain.catalog.loader import get_id_generator, get_metastore, get_warehouse
 from datachain.cli_utils import CommaSeparatedArgs
+from datachain.config import Config, ConfigLevel
 from datachain.data_storage.sqlite import (
     SQLiteDatabaseEngine,
     SQLiteIDGenerator,
@@ -26,6 +27,7 @@ from datachain.query.session import Session
 from datachain.utils import (
     ENV_DATACHAIN_GLOBAL_CONFIG_DIR,
     ENV_DATACHAIN_SYSTEM_CONFIG_DIR,
+    STUDIO_URL,
     DataChainDir,
 )
 
@@ -673,3 +675,29 @@ def dataset_rows():
         }
         for i in range(19)
     ]
+
+
+@pytest.fixture
+def studio_datasets(requests_mock):
+    with Config(ConfigLevel.GLOBAL).edit() as conf:
+        conf["studio"] = {"token": "isat_access_token", "team": "team_name"}
+
+    datasets = [
+        {
+            "id": 1,
+            "name": "dogs",
+            "versions": [{"version": 1}, {"version": 2}],
+        },
+        {
+            "id": 2,
+            "name": "cats",
+            "versions": [{"version": 1}],
+        },
+        {
+            "id": 3,
+            "name": "both",
+            "versions": [{"version": 1}],
+        },
+    ]
+
+    requests_mock.post(f"{STUDIO_URL}/api/datachain/ls-datasets", json=datasets)

{datachain-0.6.7 → datachain-0.6.8}/tests/func/test_ls.py

@@ -15,8 +15,8 @@ from datachain.lib.listing import LISTING_PREFIX
 from tests.utils import uppercase_scheme
 
 
-@pytest.fixture(autouse=True)
-def studio_config():
+@pytest.fixture
+def studio_config(global_config_dir):
     with Config(ConfigLevel.GLOBAL).edit() as conf:
         conf["studio"] = {"token": "isat_access_token", "team": "team_name"}
 
@@ -235,20 +235,11 @@ dog3
 """
 
 
-def test_ls_remote_sources(cloud_type, capsys, monkeypatch):
+def test_ls_remote_sources(cloud_type, capsys, monkeypatch, studio_config):
     src = f"{cloud_type}://bucket"
-    token = "35NmrvSlsGVxTYIglxSsBIQHRrMpi6irSSYcAL0flijOytCHc"  # noqa: S105
     with monkeypatch.context() as m:
         m.setattr("requests.post", mock_post)
-        ls(
-            [src, f"{src}/dogs/others", f"{src}/dogs"],
-            config={
-                "type": "http",
-                "url": "http://localhost:8111/api/datachain",
-                "username": "datachain-team",
-                "token": f"isat_{token}",
-            },
-        )
+        ls([src, f"{src}/dogs/others", f"{src}/dogs"], studio=True)
     captured = capsys.readouterr()
     assert captured.out == ls_remote_sources_output.format(src=src)