cognite-toolkit 0.6.97__py3-none-any.whl → 0.7.30__py3-none-any.whl

cognite_toolkit/_cdf_tk/client/api/migration.py

@@ -1,7 +1,7 @@
 import warnings
 from collections.abc import Sequence
 from itertools import groupby
-from typing import TypeVar, overload
+from typing import Literal, TypeVar, cast, overload
 
 from cognite.client._constants import DEFAULT_LIMIT_READ
 from cognite.client.data_classes.data_modeling import (
@@ -10,6 +10,7 @@ from cognite.client.data_classes.data_modeling import (
     NodeApplyResultList,
     NodeId,
     NodeList,
+    ViewId,
     filters,
     query,
 )
@@ -347,9 +348,182 @@ class SpaceSourceAPI:
         return results
 
 
+_T_Cached = TypeVar("_T_Cached", bound=NodeId | ViewId)
+
+
+class LookupAPI:
+    def __init__(self, instance_api: ExtendedInstancesAPI, resource_type: AssetCentricType) -> None:
+        self._instance_api = instance_api
+        self._resource_type = resource_type
+        self._view_id = InstanceSource.get_source()
+        self._node_id_by_id: dict[int, NodeId | None] = {}
+        self._node_id_by_external_id: dict[str, NodeId | None] = {}
+        self._consumer_view_id_by_id: dict[int, ViewId | None] = {}
+        self._consumer_view_id_by_external_id: dict[str, ViewId | None] = {}
+        self._RETRIEVE_LIMIT = 1000
+
+    @overload
+    def __call__(self, id: int, external_id: None = None) -> NodeId | None: ...
+
+    @overload
+    def __call__(self, id: SequenceNotStr[int], external_id: None = None) -> dict[int, NodeId]: ...
+
+    @overload
+    def __call__(self, *, external_id: str) -> NodeId | None: ...
+
+    @overload
+    def __call__(self, *, external_id: SequenceNotStr[str]) -> dict[str, NodeId]: ...
+
+    def __call__(
+        self, id: int | SequenceNotStr[int] | None = None, external_id: str | SequenceNotStr[str] | None = None
+    ) -> dict[int, NodeId] | dict[str, NodeId] | NodeId | None:
+        """Lookup NodeId by either internal ID or external ID.
+
+        Args:
+            id (int | Sequence[int] | None): The internal ID(s) to lookup.
+            external_id (str | SequenceNotStr[str] | None): The external ID(s) to lookup.
+
+        Returns:
+            NodeId | dict[int, NodeId] | dict[str, NodeId] | None: The corresponding NodeId(s) if found, otherwise None.
+
+        """
+        if id is not None and external_id is None:
+            return self._lookup(
+                identifier=id,
+                cache=self._node_id_by_id,
+                property_name="id",
+                return_type=NodeId,
+                input_type=int,
+            )
+        elif external_id is not None and id is None:
+            return self._lookup(
+                identifier=external_id,
+                cache=self._node_id_by_external_id,
+                property_name="classicExternalId",
+                return_type=NodeId,
+                input_type=str,
+            )
+        else:
+            raise TypeError("Either id or external_id must be provided, but not both.")
+
+    @overload
+    def consumer_view(self, id: int, external_id: None = None) -> ViewId | None: ...
+
+    @overload
+    def consumer_view(self, id: SequenceNotStr[int], external_id: None = None) -> dict[int, ViewId]: ...
+
+    @overload
+    def consumer_view(self, *, external_id: str) -> ViewId | None: ...
+
+    @overload
+    def consumer_view(self, *, external_id: SequenceNotStr[str]) -> dict[str, ViewId]: ...
+
+    def consumer_view(
+        self, id: int | SequenceNotStr[int] | None = None, external_id: str | SequenceNotStr[str] | None = None
+    ) -> dict[int, ViewId] | dict[str, ViewId] | ViewId | None:
+        """Lookup Consumer ViewId by either internal ID or external ID.
+
+        Args:
+            id (int | Sequence[int] | None): The internal ID(s) to lookup.
+            external_id (str | SequenceNotStr[str] | None): The external ID(s) to lookup.
+        Returns:
+            ViewId | dict[int, ViewId] | dict[str, ViewId] | None
+        """
+        if id is not None and external_id is None:
+            return self._lookup(
+                identifier=id,
+                cache=self._consumer_view_id_by_id,
+                property_name="id",
+                return_type=ViewId,
+                input_type=int,
+            )
+        elif external_id is not None and id is None:
+            return self._lookup(
+                identifier=external_id,
+                cache=self._consumer_view_id_by_external_id,
+                property_name="classicExternalId",
+                return_type=ViewId,
+                input_type=str,
+            )
+        else:
+            raise TypeError("Either id or external_id must be provided, but not both.")
+
+    def _lookup(
+        self,
+        identifier: _T | SequenceNotStr[_T],
+        cache: dict[_T, _T_Cached | None],
+        property_name: Literal["id", "classicExternalId"],
+        return_type: type[_T_Cached],
+        input_type: type[_T],
+    ) -> dict[_T, _T_Cached] | _T_Cached | None:
+        """Generic lookup method for both NodeId and ViewId by id or external_id."""
+        is_single = isinstance(identifier, input_type)
+        # MyPy does not understand that if is_single is True, identifier is _T, else SequenceNotStr[_T].
+        identifiers: list[_T] = [identifier] if is_single else list(identifier)  # type: ignore[arg-type, list-item]
+
+        missing = [id_ for id_ in identifiers if id_ not in cache]
+        if missing:
+            self._fetch_and_cache(missing, by=property_name)
+
+        if is_single:
+            return cache.get(identifier)  # type: ignore[arg-type]
+
+        return {id_: value for id_ in identifiers if isinstance(value := cache.get(id_), return_type)}
+
+    def _fetch_and_cache(self, identifiers: Sequence[int | str], by: Literal["id", "classicExternalId"]) -> None:
+        for chunk in chunker_sequence(identifiers, self._RETRIEVE_LIMIT):
+            retrieve_query = query.Query(
+                with_={
+                    "instanceSource": query.NodeResultSetExpression(
+                        filter=filters.And(
+                            filters.HasData(views=[self._view_id]),
+                            filters.Equals(self._view_id.as_property_ref("resourceType"), self._resource_type),
+                            filters.In(self._view_id.as_property_ref(by), list(chunk)),
+                        ),
+                        limit=len(chunk),
+                    ),
+                },
+                select={"instanceSource": query.Select([query.SourceSelector(self._view_id, ["*"])])},
+            )
+            chunk_response = self._instance_api.query(retrieve_query)
+            items = chunk_response.get("instanceSource", [])
+            for item in items:
+                instance_source = InstanceSource._load(item.dump())
+                node_id = instance_source.as_id()
+                self._node_id_by_id[instance_source.id_] = node_id
+                self._consumer_view_id_by_id[instance_source.id_] = instance_source.consumer_view()
+                if instance_source.classic_external_id:
+                    self._node_id_by_external_id[instance_source.classic_external_id] = node_id
+                    self._consumer_view_id_by_external_id[instance_source.classic_external_id] = (
+                        instance_source.consumer_view()
+                    )
+            missing = set(chunk) - set(self._node_id_by_id.keys()) - set(self._node_id_by_external_id.keys())
+            if by == "id":
+                for missing_id in cast(set[int], missing):
+                    if missing_id not in self._node_id_by_id:
+                        self._node_id_by_id[missing_id] = None
+                    if missing_id not in self._consumer_view_id_by_id:
+                        self._consumer_view_id_by_id[missing_id] = None
+            elif by == "classicExternalId":
+                for missing_ext_id in cast(set[str], missing):
+                    if missing_ext_id not in self._node_id_by_external_id:
+                        self._node_id_by_external_id[missing_ext_id] = None
+                    if missing_ext_id not in self._consumer_view_id_by_external_id:
+                        self._consumer_view_id_by_external_id[missing_ext_id] = None
+
+
+class MigrationLookupAPI:
+    def __init__(self, instance_api: ExtendedInstancesAPI) -> None:
+        self.assets = LookupAPI(instance_api, "asset")
+        self.events = LookupAPI(instance_api, "event")
+        self.files = LookupAPI(instance_api, "file")
+        self.time_series = LookupAPI(instance_api, "timeseries")
+
+
 class MigrationAPI:
     def __init__(self, instance_api: ExtendedInstancesAPI) -> None:
         self.instance_source = InstanceSourceAPI(instance_api)
         self.resource_view_mapping = ResourceViewMappingAPI(instance_api)
         self.created_source_system = CreatedSourceSystemAPI(instance_api)
         self.space_source = SpaceSourceAPI(instance_api)
+        self.lookup = MigrationLookupAPI(instance_api)

cognite_toolkit/_cdf_tk/client/api/streams.py

@@ -0,0 +1,84 @@
+from collections.abc import Sequence
+
+from rich.console import Console
+
+from cognite_toolkit._cdf_tk.client.data_classes.api_classes import PagedResponse
+from cognite_toolkit._cdf_tk.client.data_classes.streams import StreamRequest, StreamResponse
+from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, ItemsRequest, ParamRequest
+
+
+class StreamsAPI:
+    ENDPOINT = "/streams"
+
+    def __init__(self, http_client: HTTPClient, console: Console) -> None:
+        self._http_client = http_client
+        self._console = console
+        self._config = http_client.config
+
+    def create(self, items: Sequence[StreamRequest]) -> list[StreamResponse]:
+        """Create one or more streams.
+
+        Args:
+            items: Sequence of StreamRequest items to create.
+
+        Returns:
+            List of created StreamResponse items.
+        """
+        responses = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="POST",
+                items=list(items),
+            )
+        )
+        responses.raise_for_status()
+        return PagedResponse[StreamResponse].model_validate(responses.get_first_body()).items
+
+    def delete(self, external_id: str) -> None:
+        """Delete stream using its external ID.
+
+        Args:
+            external_id: External ID of the stream to delete.
+        """
+        responses = self._http_client.request_with_retries(
+            ParamRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/{external_id}"),
+                method="DELETE",
+            )
+        )
+        responses.raise_for_status()
+
+    def list(self) -> list[StreamResponse]:
+        """List streams.
+
+        Returns:
+            StreamResponseList containing the listed streams.
+        """
+        responses = self._http_client.request_with_retries(
+            ParamRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="GET",
+            )
+        )
+        responses.raise_for_status()
+        return PagedResponse[StreamResponse].model_validate(responses.get_first_body()).items
+
+    def retrieve(self, external_id: str, include_statistics: bool = True) -> StreamResponse:
+        """Retrieve a stream by its external ID.
+
+        Args:
+            external_id: External ID of the stream to retrieve.
+            include_statistics: Whether to include usage statistics in the response.
+        Returns:
+            StreamResponse item.
+        """
+        responses = self._http_client.request_with_retries(
+            ParamRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/{external_id}"),
+                method="GET",
+                parameters={"includeStatistics": include_statistics},
+            )
+        )
+        responses.raise_for_status()
+        response_body = responses.get_first_body()
+        return StreamResponse.model_validate(response_body)

cognite_toolkit/_cdf_tk/client/api/three_d.py

@@ -0,0 +1,50 @@
+from rich.console import Console
+
+from cognite_toolkit._cdf_tk.client.data_classes.api_classes import PagedResponse
+from cognite_toolkit._cdf_tk.client.data_classes.three_d import ThreeDModelResponse
+from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, ParamRequest
+from cognite_toolkit._cdf_tk.utils.useful_types import PrimitiveType
+
+
+class ThreeDModelAPI:
+    ENDPOINT = "/3d/models"
+
+    def __init__(self, http_client: HTTPClient, console: Console) -> None:
+        self._http_client = http_client
+        self._console = console
+        self._config = http_client.config
+
+    def iterate(
+        self,
+        published: bool | None = None,
+        include_revision_info: bool = False,
+        limit: int = 100,
+        cursor: str | None = None,
+    ) -> PagedResponse[ThreeDModelResponse]:
+        if not (0 < limit <= 1000):
+            raise ValueError("Limit must be between 1 and 1000.")
+        parameters: dict[str, PrimitiveType] = {
+            # There is a bug in the API. The parameter includeRevisionInfo is expected to be lower case and not
+            # camel case as documented. You get error message: Unrecognized query parameter includeRevisionInfo,
+            # did you mean includerevisioninfo?
+            "includerevisioninfo": include_revision_info,
+            "limit": limit,
+        }
+        if published is not None:
+            parameters["published"] = published
+        if cursor is not None:
+            parameters["cursor"] = cursor
+        responses = self._http_client.request_with_retries(
+            ParamRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="GET",
+                parameters=parameters,
+            )
+        )
+        responses.raise_for_status()
+        return PagedResponse[ThreeDModelResponse].model_validate(responses.get_first_body())
+
+
+class ThreeDAPI:
+    def __init__(self, http_client: HTTPClient, console: Console) -> None:
+        self.models = ThreeDModelAPI(http_client, console)

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/canvas.py

@@ -2,7 +2,7 @@ import sys
 from abc import ABC
 from collections.abc import Sequence
 from datetime import datetime, timezone
-from typing import Any
+from typing import Any, TypeVar
 from uuid import uuid4
 
 from cognite.client import CogniteClient
@@ -31,6 +31,7 @@ from cognite.client.data_classes.data_modeling.instances import (
 )
 
 from cognite_toolkit._cdf_tk.client.data_classes.migration import AssetCentricId
+from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
 
 if sys.version_info >= (3, 11):
     from typing import Self
@@ -68,6 +69,9 @@ class ExtendedTypedNodeApply(TypedNodeApply, ABC):
         return output
 
 
+T_ExtendedTypedNodeApply = TypeVar("T_ExtendedTypedNodeApply", bound=ExtendedTypedNodeApply)
+
+
 class _CanvasProperties:
     created_by = PropertyOptions("createdBy")
     updated_at = PropertyOptions("updatedAt")
@@ -901,7 +905,7 @@ class IndustrialCanvasApply(CogniteResource):
                 raise TypeError(f"Unexpected instance type: {type(instance)}")
         return ids
 
-    def dump(self, keep_existing_version: bool = True) -> dict[str, object]:
+    def dump(self, keep_existing_version: bool = True) -> dict[str, JsonVal]:
         """Dump the IndustrialCanvasApply to a dictionary."""
         return {
             "canvas": self.canvas.dump(keep_existing_version=keep_existing_version),
@@ -953,7 +957,46 @@ class IndustrialCanvasApply(CogniteResource):
     @classmethod
     def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
         """Load an IndustrialCanvasApply instance from a resource dictionary."""
-        raise NotImplementedError("IndustrialCanvasApply cannot be loaded from dict.")
+        if "canvas" not in resource:
+            raise ValueError("Resource does not contain a canvas node.")
+        canvas_resource = resource["canvas"]
+        if isinstance(canvas_resource, dict):
+            canvas = CanvasApply._load(canvas_resource)
+        elif isinstance(canvas_resource, CanvasApply):
+            canvas = canvas_resource
+        elif isinstance(canvas_resource, NodeApply):
+            canvas = CanvasApply._load(canvas_resource.dump())
+        else:
+            raise TypeError(f"Canvas resource {type(canvas_resource)} is not supported.")
+        return cls(
+            canvas=canvas,
+            annotations=cls._load_apply_items(resource.get("annotations"), CanvasAnnotationApply),
+            container_references=cls._load_apply_items(resource.get("containerReferences"), ContainerReferenceApply),
+            fdm_instance_container_references=cls._load_apply_items(
+                resource.get("fdmInstanceContainerReferences"), FdmInstanceContainerReferenceApply
+            ),
+            solution_tags=cls._load_apply_items(resource.get("solutionTags"), CogniteSolutionTagApply),
+        )
+
+    @classmethod
+    def _load_apply_items(
+        cls, items: object | None, node_cls: type[T_ExtendedTypedNodeApply]
+    ) -> list[T_ExtendedTypedNodeApply]:
+        if items is None:
+            return []
+        elif isinstance(items, Sequence):
+            nodes: list[T_ExtendedTypedNodeApply] = []
+            for node in items:
+                if isinstance(node, dict):
+                    nodes.append(node_cls._load(node))
+                elif isinstance(node, node_cls):
+                    nodes.append(node)
+                elif isinstance(node, NodeApply):
+                    nodes.append(node_cls._load(node.dump()))
+                else:
+                    raise TypeError(f"Expected a sequence of {node_cls.__name__}, got {type(node).__name__}")
+            return nodes
+        raise TypeError(f"Expected a sequence of {node_cls.__name__}, got {type(items).__name__}")
 
 
 class IndustrialCanvas(WriteableCogniteResource[IndustrialCanvasApply]):

cognite_toolkit/_cdf_tk/client/data_classes/charts.py

@@ -36,7 +36,7 @@ class ChartCore(WriteableCogniteResource["ChartWrite"], ABC):
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
         """Convert the chart to a dictionary representation."""
         output = super().dump(camel_case=camel_case)
-        output["data"] = self.data.dump(camel_case=camel_case)
+        output["data"] = self.data.model_dump(mode="json", by_alias=camel_case, exclude_unset=True)
         return output
 
 
@@ -58,7 +58,7 @@ class ChartWrite(ChartCore):
         return cls(
             external_id=resource["externalId"],
             visibility=resource["visibility"],
-            data=ChartData._load(resource["data"], cognite_client=cognite_client),
+            data=ChartData._load(resource["data"]),
         )
 
 
@@ -98,7 +98,7 @@ class Chart(ChartCore):
             created_time=resource["createdTime"],
             last_updated_time=resource["lastUpdatedTime"],
             visibility=resource["visibility"],
-            data=ChartData._load(resource["data"], cognite_client=cognite_client),
+            data=ChartData._load(resource["data"]),
             owner_id=resource["ownerId"],
         )