cognite-toolkit 0.6.111__py3-none-any.whl → 0.6.113__py3-none-any.whl

cognite_toolkit/_cdf_tk/apps/_download_app.py

@@ -2,7 +2,9 @@ from enum import Enum
 from pathlib import Path
 from typing import Annotated, Any
 
+import questionary
 import typer
+from questionary import Choice
 from rich import print
 
 from cognite_toolkit._cdf_tk.client.data_classes.raw import RawTable
@@ -11,12 +13,14 @@ from cognite_toolkit._cdf_tk.constants import DATA_DEFAULT_DIR
 from cognite_toolkit._cdf_tk.storageio import (
     AssetIO,
     ChartIO,
+    EventIO,
+    FileMetadataIO,
     HierarchyIO,
     InstanceIO,
     RawIO,
+    TimeSeriesIO,
 )
 from cognite_toolkit._cdf_tk.storageio.selectors import (
-    AssetCentricSelector,
     AssetSubtreeSelector,
     ChartExternalIdSelector,
     ChartSelector,
@@ -28,11 +32,16 @@ from cognite_toolkit._cdf_tk.storageio.selectors import (
 )
 from cognite_toolkit._cdf_tk.utils.auth import EnvironmentVariables
 from cognite_toolkit._cdf_tk.utils.interactive_select import (
+    AssetCentricInteractiveSelect,
     AssetInteractiveSelect,
     DataModelingSelect,
+    EventInteractiveSelect,
+    FileMetadataInteractiveSelect,
     InteractiveChartSelect,
     RawTableInteractiveSelect,
+    TimeSeriesInteractiveSelect,
 )
+from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricKind
 
 
 class RawFormats(str, Enum):
@@ -77,6 +86,9 @@ class DownloadApp(typer.Typer):
         self.callback(invoke_without_command=True)(self.download_main)
         self.command("raw")(self.download_raw_cmd)
         self.command("assets")(self.download_assets_cmd)
+        self.command("timeseries")(self.download_timeseries_cmd)
+        self.command("events")(self.download_events_cmd)
+        self.command("files")(self.download_files_cmd)
         self.command("hierarchy")(self.download_hierarchy_cmd)
         self.command("instances")(self.download_instances_cmd)
         self.command("charts")(self.download_charts_cmd)
@@ -178,23 +190,15 @@ class DownloadApp(typer.Typer):
             )
         )
 
-    @staticmethod
     def download_assets_cmd(
+        self,
         ctx: typer.Context,
         data_sets: Annotated[
             list[str] | None,
             typer.Option(
                 "--data-set",
                 "-d",
-                help="List of data sets to download assets from. If this and hierarchy are not provided, an interactive selection will be made.",
-            ),
-        ] = None,
-        hierarchy: Annotated[
-            list[str] | None,
-            typer.Option(
-                "--hierarchy",
-                "-r",
-                help="List of asset hierarchies to download assets from. If this and data sets are not provided, an interactive selection will be made.",
+                help="List of data sets to download assets from. If this is not provided, an interactive selection will be made.",
             ),
         ] = None,
         file_format: Annotated[
@@ -241,20 +245,17 @@ class DownloadApp(typer.Typer):
     ) -> None:
         """This command will download assets from CDF into a temporary directory."""
         client = EnvironmentVariables.create_from_environment().get_client()
-        is_interactive = not data_sets and not hierarchy
-        if is_interactive:
-            interactive = AssetInteractiveSelect(client, "download assets")
-            selector_type = interactive.select_hierarchies_or_data_sets()
-            if selector_type == "Data Set":
-                data_sets = interactive.select_data_sets()
-            else:
-                hierarchy = interactive.select_hierarchies()
-
-        selectors: list[AssetCentricSelector] = []
-        if data_sets:
-            selectors.extend([DataSetSelector(data_set_external_id=ds, kind="Assets") for ds in data_sets])
-        if hierarchy:
-            selectors.extend([AssetSubtreeSelector(hierarchy=h, kind="Assets") for h in hierarchy])
+        if data_sets is None:
+            data_sets, file_format, compression, output_dir, limit = self._asset_centric_interactive(
+                AssetInteractiveSelect(client, "download"),
+                file_format,
+                compression,
+                output_dir,
+                limit,
+                "Assets",
+            )
+
+        selectors = [DataSetSelector(kind="Assets", data_set_external_id=data_set) for data_set in data_sets]
         cmd = DownloadCommand()
         cmd.run(
             lambda: cmd.download(
@@ -268,6 +269,287 @@ class DownloadApp(typer.Typer):
             )
         )
 
+    @classmethod
+    def _asset_centric_interactive(
+        cls,
+        selector: AssetCentricInteractiveSelect,
+        file_format: AssetCentricFormats,
+        compression: CompressionFormat,
+        output_dir: Path,
+        limit: int,
+        kind: AssetCentricKind,
+    ) -> tuple[list[str], AssetCentricFormats, CompressionFormat, Path, int]:
+        data_sets = selector.select_data_sets()
+        display_name = kind.casefold() + "s"
+        file_format = questionary.select(
+            f"Select format to download the {display_name} in:",
+            choices=[Choice(title=format_.value, value=format_) for format_ in AssetCentricFormats],
+            default=file_format,
+        ).ask()
+        compression = questionary.select(
+            f"Select compression format to use when downloading the {display_name}:",
+            choices=[Choice(title=comp.value, value=comp) for comp in CompressionFormat],
+            default=compression,
+        ).ask()
+        output_dir = Path(
+            questionary.path(
+                "Where to download the assets:",
+                default=str(output_dir),
+                only_directories=True,
+            ).ask()
+        )
+        while True:
+            limit_str = questionary.text(
+                f"The maximum number of {display_name} to download from each dataset. Use -1 to download all {display_name}.",
+                default=str(limit),
+            ).ask()
+            if limit_str is None:
+                raise typer.Abort()
+            try:
+                limit = int(limit_str)
+                break
+            except ValueError:
+                print("[red]Please enter a valid integer for the limit.[/]")
+        return data_sets, file_format, compression, output_dir, limit
+
+    def download_timeseries_cmd(
+        self,
+        ctx: typer.Context,
+        data_sets: Annotated[
+            list[str] | None,
+            typer.Option(
+                "--data-set",
+                "-d",
+                help="List of data sets to download time series from. If this is not provided, an interactive selection will be made.",
+            ),
+        ] = None,
+        file_format: Annotated[
+            AssetCentricFormats,
+            typer.Option(
+                "--format",
+                "-f",
+                help="Format to download the time series in.",
+            ),
+        ] = AssetCentricFormats.csv,
+        compression: Annotated[
+            CompressionFormat,
+            typer.Option(
+                "--compression",
+                "-z",
+                help="Compression format to use when downloading the time series.",
+            ),
+        ] = CompressionFormat.none,
+        output_dir: Annotated[
+            Path,
+            typer.Option(
+                "--output-dir",
+                "-o",
+                help="Where to download the time series.",
+                allow_dash=True,
+            ),
+        ] = DEFAULT_DOWNLOAD_DIR,
+        limit: Annotated[
+            int,
+            typer.Option(
+                "--limit",
+                "-l",
+                help="The maximum number of time series to download from each dataset. Use -1 to download all time series.",
+            ),
+        ] = 100_000,
+        verbose: Annotated[
+            bool,
+            typer.Option(
+                "--verbose",
+                "-v",
+                help="Turn on to get more verbose output when running the command",
+            ),
+        ] = False,
+    ) -> None:
+        """This command will download time series from CDF into a temporary directory."""
+        client = EnvironmentVariables.create_from_environment().get_client()
+        if data_sets is None:
+            data_sets, file_format, compression, output_dir, limit = self._asset_centric_interactive(
+                TimeSeriesInteractiveSelect(client, "download"),
+                file_format,
+                compression,
+                output_dir,
+                limit,
+                "TimeSeries",
+            )
+
+        selectors = [DataSetSelector(kind="TimeSeries", data_set_external_id=data_set) for data_set in data_sets]
+        cmd = DownloadCommand()
+        cmd.run(
+            lambda: cmd.download(
+                selectors=selectors,
+                io=TimeSeriesIO(client),
+                output_dir=output_dir,
+                file_format=f".{file_format.value}",
+                compression=compression.value,
+                limit=limit if limit != -1 else None,
+                verbose=verbose,
+            )
+        )
+
+    def download_events_cmd(
+        self,
+        ctx: typer.Context,
+        data_sets: Annotated[
+            list[str] | None,
+            typer.Option(
+                "--data-set",
+                "-d",
+                help="List of data sets to download events from. If this is not provided, an interactive selection will be made.",
+            ),
+        ] = None,
+        file_format: Annotated[
+            AssetCentricFormats,
+            typer.Option(
+                "--format",
+                "-f",
+                help="Format to download the events in.",
+            ),
+        ] = AssetCentricFormats.csv,
+        compression: Annotated[
+            CompressionFormat,
+            typer.Option(
+                "--compression",
+                "-z",
+                help="Compression format to use when downloading the events.",
+            ),
+        ] = CompressionFormat.none,
+        output_dir: Annotated[
+            Path,
+            typer.Option(
+                "--output-dir",
+                "-o",
+                help="Where to download the events.",
+                allow_dash=True,
+            ),
+        ] = DEFAULT_DOWNLOAD_DIR,
+        limit: Annotated[
+            int,
+            typer.Option(
+                "--limit",
+                "-l",
+                help="The maximum number of events to download from each dataset. Use -1 to download all events.",
+            ),
+        ] = 100_000,
+        verbose: Annotated[
+            bool,
+            typer.Option(
+                "--verbose",
+                "-v",
+                help="Turn on to get more verbose output when running the command",
+            ),
+        ] = False,
+    ) -> None:
+        """This command will download events from CDF into a temporary directory."""
+        client = EnvironmentVariables.create_from_environment().get_client()
+        if data_sets is None:
+            data_sets, file_format, compression, output_dir, limit = self._asset_centric_interactive(
+                EventInteractiveSelect(client, "download"),
+                file_format,
+                compression,
+                output_dir,
+                limit,
+                "Events",
+            )
+
+        selectors = [DataSetSelector(kind="Events", data_set_external_id=data_set) for data_set in data_sets]
+        cmd = DownloadCommand()
+
+        cmd.run(
+            lambda: cmd.download(
+                selectors=selectors,
+                io=EventIO(client),
+                output_dir=output_dir,
+                file_format=f".{file_format.value}",
+                compression=compression.value,
+                limit=limit if limit != -1 else None,
+                verbose=verbose,
+            )
+        )
+
+    def download_files_cmd(
+        self,
+        ctx: typer.Context,
+        data_sets: Annotated[
+            list[str] | None,
+            typer.Option(
+                "--data-set",
+                "-d",
+                help="List of data sets to download file metadata from. If this is not provided, an interactive selection will be made.",
+            ),
+        ] = None,
+        file_format: Annotated[
+            AssetCentricFormats,
+            typer.Option(
+                "--format",
+                "-f",
+                help="Format to download the file metadata in.",
+            ),
+        ] = AssetCentricFormats.csv,
+        compression: Annotated[
+            CompressionFormat,
+            typer.Option(
+                "--compression",
+                "-z",
+                help="Compression format to use when downloading the file metadata.",
+            ),
+        ] = CompressionFormat.none,
+        output_dir: Annotated[
+            Path,
+            typer.Option(
+                "--output-dir",
+                "-o",
+                help="Where to download the file metadata.",
+                allow_dash=True,
+            ),
+        ] = DEFAULT_DOWNLOAD_DIR,
+        limit: Annotated[
+            int,
+            typer.Option(
+                "--limit",
+                "-l",
+                help="The maximum number of file metadata to download from each dataset. Use -1 to download all file metadata.",
+            ),
+        ] = 100_000,
+        verbose: Annotated[
+            bool,
+            typer.Option(
+                "--verbose",
+                "-v",
+                help="Turn on to get more verbose output when running the command",
+            ),
+        ] = False,
+    ) -> None:
+        """This command will download file metadata from CDF into a temporary directory."""
+        client = EnvironmentVariables.create_from_environment().get_client()
+        if data_sets is None:
+            data_sets, file_format, compression, output_dir, limit = self._asset_centric_interactive(
+                FileMetadataInteractiveSelect(client, "download"),
+                file_format,
+                compression,
+                output_dir,
+                limit,
+                "FileMetadata",
+            )
+
+        selectors = [DataSetSelector(kind="FileMetadata", data_set_external_id=data_set) for data_set in data_sets]
+        cmd = DownloadCommand()
+        cmd.run(
+            lambda: cmd.download(
+                selectors=selectors,
+                io=FileMetadataIO(client),
+                output_dir=output_dir,
+                file_format=f".{file_format.value}",
+                compression=compression.value,
+                limit=limit if limit != -1 else None,
+                verbose=verbose,
+            )
+        )
+
     @staticmethod
     def download_hierarchy_cmd(
         ctx: typer.Context,

cognite_toolkit/_cdf_tk/client/data_classes/base.py

@@ -1,10 +1,14 @@
 import sys
 from abc import ABC, abstractmethod
-from typing import Any, Generic, TypeVar
+from collections import UserList
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 from pydantic import BaseModel, ConfigDict
 from pydantic.alias_generators import to_camel
 
+if TYPE_CHECKING:
+    from cognite.client import CogniteClient
+
 if sys.version_info >= (3, 11):
     from typing import Self
 else:
@@ -61,3 +65,23 @@ class Identifier(BaseModel):
 
     def as_id(self) -> Self:
         return self
+
+
+T_Resource = TypeVar("T_Resource", bound=RequestResource | ResponseResource)
+
+
+class BaseResourceList(UserList[T_Resource]):
+    """Base class for resource lists."""
+
+    _RESOURCE: type[T_Resource]
+
+    def __init__(self, initlist: list[T_Resource] | None = None, **_: Any) -> None:
+        super().__init__(initlist or [])
+
+    def dump(self, camel_case: bool = True) -> list[dict[str, Any]]:
+        return [item.dump(camel_case) for item in self.data]
+
+    @classmethod
+    def load(cls, data: list[dict[str, Any]], cognite_client: "CogniteClient | None" = None) -> Self:
+        items = [cls._RESOURCE.model_validate(item) for item in data]
+        return cls(items)  # type: ignore[arg-type]

cognite_toolkit/_cdf_tk/client/data_classes/infield.py

@@ -1,15 +1,16 @@
 import sys
-from collections import UserList
 from typing import Any, ClassVar, Literal
 
-from cognite.client import CogniteClient
 from pydantic import JsonValue, field_validator
 from pydantic_core.core_schema import ValidationInfo
 
-from cognite_toolkit._cdf_tk.protocols import ResourceRequestListProtocol, ResourceResponseListProtocol
+from cognite_toolkit._cdf_tk.protocols import (
+    ResourceRequestListProtocol,
+    ResourceResponseListProtocol,
+)
 from cognite_toolkit._cdf_tk.utils.text import sanitize_instance_external_id
 
-from .base import ResponseResource
+from .base import BaseResourceList, ResponseResource
 from .instance_api import InstanceRequestResource, ViewReference
 
 if sys.version_info >= (3, 11):
@@ -74,29 +75,13 @@ class InfieldLocationConfig(
 
 
 class InfieldLocationConfigList(
-    UserList[InfieldLocationConfig],
+    BaseResourceList[InfieldLocationConfig],
     ResourceResponseListProtocol,
     ResourceRequestListProtocol,
 ):
     """A list of InfieldLocationConfig objects."""
 
     _RESOURCE = InfieldLocationConfig
-    data: list[InfieldLocationConfig]
-
-    def __init__(self, initlist: list[InfieldLocationConfig] | None = None, **_: Any) -> None:
-        super().__init__(initlist or [])
-
-    def dump(self, camel_case: bool = True) -> list[dict[str, Any]]:
-        """Serialize the list of InfieldLocationConfig objects to a list of dictionaries."""
-        return [item.dump(camel_case) for item in self.data]
-
-    @classmethod
-    def load(
-        cls, data: list[dict[str, Any]], cognite_client: CogniteClient | None = None
-    ) -> "InfieldLocationConfigList":
-        """Deserialize a list of dictionaries to an InfieldLocationConfigList."""
-        items = [InfieldLocationConfig.model_validate(item) for item in data]
-        return cls(items)
 
     def as_write(self) -> Self:
         return self

cognite_toolkit/_cdf_tk/client/data_classes/streams.py

@@ -1,18 +1,12 @@
-import sys
-from collections import UserList
-from typing import Any, Literal
-
-from cognite.client import CogniteClient
+from typing import Literal
 
 from cognite_toolkit._cdf_tk.constants import StreamTemplateName
-from cognite_toolkit._cdf_tk.protocols import ResourceRequestListProtocol, ResourceResponseListProtocol
-
-from .base import BaseModelObject, RequestResource, ResponseResource
+from cognite_toolkit._cdf_tk.protocols import (
+    ResourceRequestListProtocol,
+    ResourceResponseListProtocol,
+)
 
-if sys.version_info >= (3, 11):
-    from typing import Self
-else:
-    from typing_extensions import Self
+from .base import BaseModelObject, BaseResourceList, RequestResource, ResponseResource
 
 
 class StreamRequest(RequestResource):
@@ -25,22 +19,11 @@ class StreamRequest(RequestResource):
         return self.external_id
 
 
-class StreamRequestList(UserList[StreamRequest], ResourceRequestListProtocol):
+class StreamRequestList(BaseResourceList[StreamRequest], ResourceRequestListProtocol):
     """List of Stream request resources."""
 
     _RESOURCE = StreamRequest
 
-    def __init__(self, initlist: list[StreamRequest] | None = None, **_: Any) -> None:
-        super().__init__(initlist or [])
-
-    def dump(self, camel_case: bool = True) -> list[dict[str, Any]]:
-        return [item.dump(camel_case) for item in self.data]
-
-    @classmethod
-    def load(cls, data: list[dict[str, Any]], cognite_client: CogniteClient | None = None) -> "StreamRequestList":
-        items = [StreamRequest.model_validate(item) for item in data]
-        return cls(items)
-
 
 class LifecycleObject(BaseModelObject):
     """Lifecycle object."""
@@ -98,22 +81,10 @@ class StreamResponse(ResponseResource["StreamRequest"]):
         )
 
 
-class StreamResponseList(UserList[StreamResponse], ResourceResponseListProtocol):
+class StreamResponseList(BaseResourceList[StreamResponse], ResourceResponseListProtocol):
     """List of Stream response resources."""
 
     _RESOURCE = StreamResponse
-    data: list[StreamResponse]
-
-    def __init__(self, initlist: list[StreamResponse] | None = None, **_: Any) -> None:
-        super().__init__(initlist or [])
-
-    def dump(self, camel_case: bool = True) -> list[dict[str, Any]]:
-        return [item.dump(camel_case) for item in self.data]
-
-    @classmethod
-    def load(cls, data: list[dict[str, Any]], cognite_client: CogniteClient | None = None) -> "StreamResponseList":
-        items = [StreamResponse.model_validate(item) for item in data]
-        return cls(items)
 
-    def as_write(self) -> Self:
-        return self
+    def as_write(self) -> StreamRequestList:
+        return StreamRequestList([item.as_write() for item in self.data])

cognite_toolkit/_cdf_tk/commands/_upload.py

@@ -11,13 +11,12 @@ from pydantic import ValidationError
 from rich.console import Console
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.constants import DATA_MANIFEST_STEM, DATA_RESOURCE_DIR
+from cognite_toolkit._cdf_tk.constants import DATA_MANIFEST_SUFFIX, DATA_RESOURCE_DIR
 from cognite_toolkit._cdf_tk.cruds import ViewCRUD
 from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
 from cognite_toolkit._cdf_tk.storageio import (
     T_Selector,
     UploadableStorageIO,
-    are_same_kind,
     get_upload_io,
 )
 from cognite_toolkit._cdf_tk.storageio._base import T_WriteCogniteResource, TableUploadableStorageIO, UploadItem
@@ -88,7 +87,7 @@ class UploadCommand(ToolkitCommand):
         └── ...
         """
         console = client.console
-        data_files_by_selector = self._find_data_files(input_dir, kind)
+        data_files_by_selector = self._find_data_files(input_dir)
 
         self._deploy_resource_folder(input_dir / DATA_RESOURCE_DIR, deploy_resources, client, console, dry_run, verbose)
 
@@ -138,38 +137,26 @@ class UploadCommand(ToolkitCommand):
     def _find_data_files(
         self,
         input_dir: Path,
-        kind: str | None = None,
     ) -> dict[Selector, list[Path]]:
         """Finds data files and their corresponding metadata files in the input directory."""
-        manifest_file_endswith = f".{DATA_MANIFEST_STEM}.yaml"
         data_files_by_metadata: dict[Selector, list[Path]] = {}
-        for metadata_file in input_dir.glob(f"*{manifest_file_endswith}"):
-            data_file_prefix = metadata_file.name.removesuffix(manifest_file_endswith)
-            data_files = [
-                file
-                for file in input_dir.glob(f"{data_file_prefix}*")
-                if not file.name.endswith(manifest_file_endswith)
-            ]
-            if kind is not None and data_files:
-                data_files = [data_file for data_file in data_files if are_same_kind(kind, data_file)]
-                if not data_files:
-                    continue
-            if not data_files:
-                self.warn(
-                    MediumSeverityWarning(
-                        f"Metadata file {metadata_file.as_posix()!r} has no corresponding data files, skipping.",
-                    )
-                )
-                continue
-
-            selector_dict = read_yaml_file(metadata_file, expected_output="dict")
+        for manifest_file in input_dir.glob(f"*{DATA_MANIFEST_SUFFIX}"):
+            selector_dict = read_yaml_file(manifest_file, expected_output="dict")
             try:
                 selector = SelectorAdapter.validate_python(selector_dict)
             except ValidationError as e:
                 errors = humanize_validation_error(e)
                 self.warn(
                     ResourceFormatWarning(
-                        metadata_file, tuple(errors), text="Invalid selector in metadata file, skipping."
+                        manifest_file, tuple(errors), text="Invalid selector in metadata file, skipping."
+                    )
+                )
+                continue
+            data_files = selector.find_data_files(input_dir, manifest_file)
+            if not data_files:
+                self.warn(
+                    MediumSeverityWarning(
+                        f"Metadata file {manifest_file.as_posix()!r} has no corresponding data files, skipping.",
                     )
                 )
                 continue
@@ -230,7 +217,7 @@ class UploadCommand(ToolkitCommand):
                     raise ToolkitValueError(f"{selector.display_name} does not support {reader.format!r} files.")
                 tracker = ProgressTracker[str]([self._UPLOAD])
                 executor = ProducerWorkerExecutor[list[tuple[str, dict[str, JsonVal]]], Sequence[UploadItem]](
-                    download_iterable=io.read_chunks(reader),
+                    download_iterable=io.read_chunks(reader, selector),
                     process=partial(io.rows_to_data, selector=selector)
                     if reader.is_table and isinstance(io, TableUploadableStorageIO)
                     else io.json_chunk_to_data,
@@ -275,7 +262,7 @@ class UploadCommand(ToolkitCommand):
         self, selector: Selector, data_file: Path, client: ToolkitClient
     ) -> UploadableStorageIO | None:
         try:
-            io_cls = get_upload_io(type(selector), kind=data_file)
+            io_cls = get_upload_io(type(selector))
         except ValueError as e:
             self.warn(HighSeverityWarning(f"Could not find StorageIO for selector {selector}: {e}"))
             return None