cognite-toolkit 0.5.97__py3-none-any.whl → 0.7.34__py3-none-any.whl

cognite_toolkit/_cdf_tk/utils/aggregators.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from functools import lru_cache
 from typing import ClassVar, Generic, Literal, TypeVar
 
 from cognite.client.data_classes import (
@@ -17,6 +18,7 @@ from cognite.client.data_classes.sequences import SequenceProperty
 from cognite.client.data_classes.time_series import TimeSeriesProperty
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.exceptions import ToolkitMissingResourceError
 from cognite_toolkit._cdf_tk.utils.cdf import (
     label_aggregate_count,
     label_count,
@@ -96,7 +98,7 @@ class AssetCentricAggregator(ABC):
                 ids.append(id_)
                 seen.add(id_)
             try:
-                int_id = int(id_)  # type: ignore[arg-type]
+                int_id = int(id_)
             except (ValueError, TypeError):
                 continue
             if int_id not in seen:
@@ -104,6 +106,15 @@ class AssetCentricAggregator(ABC):
                 seen.add(int_id)
         return ids
 
+    def _to_dataset_id(self, data_set_external_id: str | list[str] | None) -> list[int] | None:
+        """Converts data set external IDs to data set IDs."""
+        dataset_id: list[int] | None = None
+        if data_set_external_id is not None:
+            if isinstance(data_set_external_id, str):
+                data_set_external_id = [data_set_external_id]
+            dataset_id = self.client.lookup.data_sets.id(data_set_external_id, allow_empty=False)
+        return dataset_id
+
 
 class MetadataAggregator(AssetCentricAggregator, ABC, Generic[T_CogniteFilter]):
     filter_cls: type[T_CogniteFilter]
@@ -114,12 +125,79 @@ class MetadataAggregator(AssetCentricAggregator, ABC, Generic[T_CogniteFilter]):
         super().__init__(client)
         self.resource_name = resource_name
 
-    def metadata_key_count(self) -> int:
-        return len(metadata_key_counts(self.client, self.resource_name))
+    def metadata_key_count(
+        self, hierarchy: str | list[str] | None = None, data_sets: str | list[str] | None = None
+    ) -> int:
+        """Returns the number of metadata keys used by the resource."""
+        return len(self.used_metadata_keys(hierarchy=hierarchy, data_sets=data_sets))
+
+    def used_metadata_keys(
+        self, hierarchy: str | list[str] | None = None, data_sets: str | list[str] | None = None
+    ) -> list[tuple[str, int]]:
+        """Returns a list of metadata keys and their counts."""
+        hierarchy_ids, data_set_ids = self._lookup_hierarchy_data_set_pair(
+            hierarchy, data_sets, operation="find metadata keys"
+        )
+        return self._used_metadata_keys(hierarchy=hierarchy_ids, data_sets=data_set_ids)
+
+    @lru_cache(maxsize=1)
+    def _used_metadata_keys(
+        self, hierarchy: tuple[int, ...] | None = None, data_sets: tuple[int, ...] | None = None
+    ) -> list[tuple[str, int]]:
+        return metadata_key_counts(
+            self.client,
+            self.resource_name,
+            hierarchies=list(hierarchy) if hierarchy else None,
+            data_sets=list(data_sets) if data_sets else None,
+        )
+
+    def _lookup_hierarchy_data_set_pair(
+        self, hierarchy: str | list[str] | None, data_sets: str | list[str] | None, operation: str
+    ) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None]:
+        """Returns a tuple of hierarchy and data sets."""
+        hierarchy_ids: tuple[int, ...] | None = None
+        if isinstance(hierarchy, str):
+            asset_id = self.client.lookup.assets.id(external_id=hierarchy, allow_empty=False)
+            if asset_id is None:
+                raise ToolkitMissingResourceError(f"Cannot {operation}. Asset with external ID {hierarchy!r} not found")
+            hierarchy_ids = (asset_id,)
+        elif isinstance(hierarchy, list) and all(isinstance(item, str) for item in hierarchy):
+            asset_ids = self.client.lookup.assets.id(external_id=hierarchy, allow_empty=False)
+            if len(asset_ids) != len(hierarchy):
+                missing = set(hierarchy) - set(
+                    self.client.lookup.assets.external_id([id_ for id_ in asset_ids if id_ is not None])
+                )
+                raise ToolkitMissingResourceError(
+                    f"Cannot {operation}. Assets with external IDs {sorted(missing)!r} not found"
+                )
+            hierarchy_ids = tuple(sorted(asset_ids))
+
+        data_set_ids: tuple[int, ...] | None = None
+        if isinstance(data_sets, str):
+            data_set_id = self.client.lookup.data_sets.id(external_id=data_sets, allow_empty=False)
+            if data_set_id is None:
+                raise ToolkitMissingResourceError(
+                    f"Cannot {operation}. Data set with external ID {data_sets!r} not found"
+                )
+            data_set_ids = (data_set_id,)
+        elif isinstance(data_sets, list) and all(isinstance(item, str) for item in data_sets):
+            data_set_ids_list = self.client.lookup.data_sets.id(external_id=data_sets, allow_empty=False)
+            if len(data_set_ids_list) != len(data_sets):
+                missing = set(data_sets) - set(
+                    self.client.lookup.data_sets.external_id([id_ for id_ in data_set_ids_list if id_ is not None])
+                )
+                raise ToolkitMissingResourceError(
+                    f"Cannot {operation}. Data sets with external IDs {sorted(missing)!r} not found"
+                )
+            data_set_ids = tuple(sorted(data_set_ids_list))
+
+        return hierarchy_ids, data_set_ids
 
     @classmethod
     def create_filter(
-        cls, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
+        cls,
+        hierarchy: str | list[str] | tuple[str, ...] | None = None,
+        data_set_external_id: str | list[str] | tuple[str, ...] | None = None,
     ) -> T_CogniteFilter | None:
         """Creates a filter for the resource based on hierarchy and data set external ID."""
         if cls._is_empty(hierarchy) and cls._is_empty(data_set_external_id):
@@ -127,29 +205,52 @@ class MetadataAggregator(AssetCentricAggregator, ABC, Generic[T_CogniteFilter]):
         asset_subtree_ids: list[dict[str, str]] | None = None
         if isinstance(hierarchy, str):
             asset_subtree_ids = [{"externalId": hierarchy}]
-        elif isinstance(hierarchy, list) and hierarchy:
+        elif isinstance(hierarchy, list | tuple) and hierarchy:
             asset_subtree_ids = [{"externalId": item} for item in hierarchy]
         data_set_ids: list[dict[str, str]] | None = None
         if isinstance(data_set_external_id, str):
             data_set_ids = [{"externalId": data_set_external_id}]
-        elif isinstance(data_set_external_id, list) and data_set_external_id:
+        elif isinstance(data_set_external_id, list | tuple) and data_set_external_id:
             data_set_ids = [{"externalId": item} for item in data_set_external_id]
 
-        return cls.filter_cls(asset_subtree_ids=asset_subtree_ids, data_set_ids=data_set_ids)
+        # MyPy fails to understand that filter_cls() produce a T_CogniteFilter
+        return cls.filter_cls(asset_subtree_ids=asset_subtree_ids, data_set_ids=data_set_ids)  # type: ignore[return-value]
 
     @classmethod
-    def _is_empty(cls, items: str | list[str] | None) -> bool:
+    def _is_empty(cls, items: str | list[str] | tuple[str, ...] | None) -> bool:
         """Checks if the provided items are empty."""
         if items is None:
             return True
-        if isinstance(items, list):
+        if isinstance(items, list | tuple):
             return not items
         return False
 
 
 class LabelAggregator(MetadataAggregator, ABC, Generic[T_CogniteFilter]):
-    def label_count(self) -> int:
-        return len(label_count(self.client, self.resource_name))
+    def label_count(self, hierarchy: str | list[str] | None = None, data_sets: str | list[str] | None = None) -> int:
+        """Returns the number of labels used by the resource."""
+        return len(self.used_labels(hierarchy=hierarchy, data_sets=data_sets))
+
+    def used_labels(
+        self, hierarchy: str | list[str] | None = None, data_sets: str | list[str] | None = None
+    ) -> list[tuple[str, int]]:
+        """Returns a list of labels and their counts."""
+        hierarchy_ids, data_set_ids = self._lookup_hierarchy_data_set_pair(
+            hierarchy, data_sets, operation="find labels"
+        )
+        return self._used_labels(hierarchy=hierarchy_ids, data_sets=data_set_ids)
+
+    @lru_cache(maxsize=1)
+    def _used_labels(
+        self, hierarchy: tuple[int, ...] | None = None, data_sets: tuple[int, ...] | None = None
+    ) -> list[tuple[str, int]]:
+        """Returns a list of labels and their counts."""
+        return label_count(
+            self.client,
+            self.resource_name,
+            hierarchies=list(hierarchy) if hierarchy else None,
+            data_sets=list(data_sets) if data_sets else None,
+        )
 
 
 class AssetAggregator(LabelAggregator[AssetFilter]):
@@ -164,7 +265,9 @@ class AssetAggregator(LabelAggregator[AssetFilter]):
         return "Assets"
 
     def count(
-        self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
+        self,
+        hierarchy: str | list[str] | tuple[str, ...] | None = None,
+        data_set_external_id: str | list[str] | tuple[str, ...] | None = None,
     ) -> int:
         return self.client.assets.aggregate_count(filter=self.create_filter(hierarchy, data_set_external_id))
 
@@ -292,9 +395,10 @@ class RelationshipAggregator(AssetCentricAggregator):
     def count(
         self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
     ) -> int:
-        if hierarchy is not None or data_set_external_id is not None:
+        if hierarchy is not None:
             raise NotImplementedError()
-        results = relationship_aggregate_count(self.client)
+        dataset_id = self._to_dataset_id(data_set_external_id)
+        results = relationship_aggregate_count(self.client, dataset_id)
         return sum(result.count for result in results)
 
     def used_data_sets(self, hierarchy: str | None = None) -> list[str]:
@@ -311,9 +415,10 @@ class LabelCountAggregator(AssetCentricAggregator):
     def count(
         self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
     ) -> int:
-        if hierarchy is not None or data_set_external_id is not None:
+        if hierarchy is not None:
             raise NotImplementedError()
-        return label_aggregate_count(self.client)
+        data_set_id = self._to_dataset_id(data_set_external_id)
+        return label_aggregate_count(self.client, data_set_id)
 
     def used_data_sets(self, hierarchy: str | None = None) -> list[str]:
         raise NotImplementedError()

cognite_toolkit/_cdf_tk/utils/auxiliary.py

@@ -0,0 +1,23 @@
+import functools
+import platform
+
+from cognite_toolkit import _version
+
+
+def get_current_toolkit_version() -> str:
+    return _version.__version__
+
+
+@functools.lru_cache(maxsize=1)
+def get_user_agent() -> str:
+    toolkit_version = f"CogniteToolkit/{get_current_toolkit_version()}"
+    python_version = (
+        f"{platform.python_implementation()}/{platform.python_version()} "
+        f"({platform.python_build()};{platform.python_compiler()})"
+    )
+    os_version_info = [platform.release(), platform.machine(), platform.architecture()[0]]
+    os_version_info = [s for s in os_version_info if s]  # Ignore empty strings
+    os_version_info_str = "-".join(os_version_info)
+    operating_system = f"{platform.system()}/{os_version_info_str}"
+
+    return f"{toolkit_version} {python_version} {operating_system}"

cognite_toolkit/_cdf_tk/utils/cdf.py

@@ -36,7 +36,7 @@ from cognite_toolkit._cdf_tk.tk_warnings import (
 )
 from cognite_toolkit._cdf_tk.utils import humanize_collection
 
-from .file import to_directory_compatible
+from .file import sanitize_filename
 from .sql_parser import SQLParser
 
 if sys.version_info < (3, 11):
@@ -103,10 +103,14 @@ def iterate_instances(
     # By making the sort be on external_id, Postgres should pick the index that's on (project_id, space, external_id)
     # WHERE deleted_at IS NULL. In other words, avoiding soft deleted instances.
     body["sort"] = [
+        {
+            "property": [instance_type, "space"],
+            "direction": "ascending",
+        },
         {
             "property": [instance_type, "externalId"],
             "direction": "ascending",
-        }
+        },
     ]
     url = f"/api/{client._API_VERSION}/projects/{client.config.project}/models/instances/list"
     if space:
@@ -137,7 +141,7 @@ def iterate_instances(
 
 @overload
 def read_auth(
-    authentication: object,
+    authentication: Any,
     client_config: ToolkitClientConfig,
     identifier: Hashable,
     resource_name: str,
@@ -148,7 +152,7 @@ def read_auth(
 
 @overload
 def read_auth(
-    authentication: object,
+    authentication: Any,
     client_config: ToolkitClientConfig,
     identifier: Hashable,
     resource_name: str,
@@ -158,7 +162,7 @@ def read_auth(
 
 
 def read_auth(
-    authentication: object,
+    authentication: Any,
     client_config: ToolkitClientConfig,
     identifier: Hashable,
     resource_name: str,
@@ -406,7 +410,7 @@ class ThrottlerState:
     @classmethod
     def _filepath(cls, project: str) -> Path:
         # Enure the project name is a valid string for hashing.
-        project = to_directory_compatible(project)
+        project = sanitize_filename(project)
         filename = cls._FILENAME.format(project=project)
         return Path(tempfile.gettempdir()) / filename
 

cognite_toolkit/_cdf_tk/utils/cli_args.py

@@ -0,0 +1,34 @@
+from cognite.client.data_classes.data_modeling import ViewId
+
+from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
+
+
+def parse_view_str(view_str: str) -> ViewId:
+    """Parse a view string into a ViewId.
+
+    Args:
+        view_str (str): The view string to parse.
+
+    Returns:
+        ViewId: The parsed ViewId.
+
+    Raises:
+        ToolkitValueError: If the view string is not in a valid format.
+
+    >>> parse_view_str("my_space:my_view/v1")
+    ViewId(space='my_space', external_id='my_view', version='v1')
+    """
+    space_and_rest = view_str.split(":", 1)
+    if len(space_and_rest) != 2 or not space_and_rest[0] or not space_and_rest[1]:
+        raise ToolkitValueError(
+            f"Invalid view string format: '{view_str}'. Expected format 'space:externalId/version'."
+        )
+    space = space_and_rest[0]
+    external_id_and_version = space_and_rest[1].rsplit("/", 1)
+    if len(external_id_and_version) != 2 or not external_id_and_version[0] or not external_id_and_version[1]:
+        raise ToolkitValueError(
+            f"Invalid view string format: '{view_str}'. Expected format 'space:externalId/version'."
+        )
+    external_id = external_id_and_version[0]
+    version = external_id_and_version[1]
+    return ViewId(space=space, external_id=external_id, version=version)

cognite_toolkit/_cdf_tk/utils/collection.py

@@ -1,5 +1,5 @@
 import difflib
-from collections.abc import Collection, Iterable, Iterator, Sequence
+from collections.abc import Collection, Iterable, Iterator, Sequence, Set
 from itertools import islice
 from typing import Any, TypeVar
 
@@ -18,6 +18,47 @@ def flatten_dict(dct: dict[str, Any]) -> dict[tuple[str, ...], Any]:
     return items
 
 
+def flatten_dict_json_path(dct: dict[str, Any], keep_structured: Set[str] | None = None) -> dict[str, Any]:
+    """Flatten a dictionary to a dictionary with JSON path keys.
+
+    Empty keys are ignored.
+
+    Args:
+        dct: The dictionary to flatten.
+        keep_structured: A set of keys to keep structured (not flatten). If a key is in this set,
+            it will not be flattened, and the value will be kept as is. This only applies to top-level keys.
+
+    Returns:
+        A dictionary with JSON path keys.
+    """
+
+    return _flatten(dct, keep_structured=keep_structured or set())
+
+
+def _flatten(obj: Any, keep_structured: Set[str], path: str = "") -> dict[str, Any]:
+    items: dict[str, Any] = {}
+    if isinstance(obj, dict):
+        for key, value in obj.items():
+            if not key:
+                continue
+            current_path = f"{path}.{key}" if path else key
+            if key in keep_structured:
+                items[current_path] = value
+            else:
+                items.update(_flatten(value, set(), current_path))
+    elif isinstance(obj, list):
+        for i, value in enumerate(obj):
+            current_path = f"{path}[{i}]"
+            if isinstance(value, (dict, list)):
+                items.update(_flatten(value, set(), current_path))
+            else:
+                items[current_path] = value
+    else:
+        items[path] = obj
+
+    return items
+
+
 def to_diff(a: dict[str, Any], b: dict[str, Any]) -> Iterator[str]:
     a_str = yaml_safe_dump(a, sort_keys=True)
     b_str = yaml_safe_dump(b, sort_keys=True)