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

cognite_toolkit/_cdf_tk/client/data_classes/base.py

@@ -0,0 +1,87 @@
+import sys
+from abc import ABC, abstractmethod
+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:
+    from typing_extensions import Self
+
+
+class BaseModelObject(BaseModel):
+    """Base class for all object. This includes resources and nested objects."""
+
+    # We allow extra fields to support forward compatibility.
+    model_config = ConfigDict(alias_generator=to_camel, extra="allow")
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        """Dump the resource to a dictionary.
+
+        This is the default serialization method for request resources.
+        """
+        return self.model_dump(mode="json", by_alias=camel_case)
+
+    @classmethod
+    def _load(cls, resource: dict[str, Any]) -> "Self":
+        """Load method to match CogniteResource signature."""
+        return cls.model_validate(resource)
+
+
+class RequestResource(BaseModelObject): ...
+
+
+T_RequestResource = TypeVar("T_RequestResource", bound=RequestResource)
+
+
+class ResponseResource(BaseModelObject, Generic[T_RequestResource], ABC):
+    @abstractmethod
+    def as_request_resource(self) -> T_RequestResource:
+        """Convert the response resource to a request resource."""
+        ...
+
+
+class Identifier(BaseModel):
+    """Base class for all identifier classes."""
+
+    model_config = ConfigDict(alias_generator=to_camel, extra="ignore", populate_by_name=True, frozen=True)
+
+    def dump(self, include_type: bool = True) -> dict[str, Any]:
+        """Dump the identifier to a dictionary.
+
+        Args:
+            include_type (bool): Whether to include the type of the identifier in the output.
+
+        Returns:
+            dict[str, Any]: The dumped identifier.
+        """
+        return self.model_dump(mode="json", by_alias=True, exclude_defaults=not include_type)
+
+    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

@@ -1,9 +1,17 @@
+import sys
 from abc import ABC
-from collections.abc import Mapping, Sequence
+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
+from cognite.client.data_classes._base import (
+    CogniteResource,
+    CogniteResourceList,
+    WriteableCogniteResource,
+    WriteableCogniteResourceList,
+)
 from cognite.client.data_classes.data_modeling import (
     DirectRelationReference,
     EdgeId,
@@ -23,6 +31,13 @@ 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
+else:
+    from typing_extensions import Self
+
 
 CANVAS_INSTANCE_SPACE = "IndustrialCanvasInstanceSpace"
 SOLUTION_TAG_SPACE = "SolutionTagsInstanceSpace"
@@ -54,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")
@@ -815,7 +833,7 @@ class FdmInstanceContainerReference(_FdmInstanceContainerReferenceProperties, Ty
         )
 
 
-class IndustrialCanvasApply:
+class IndustrialCanvasApply(CogniteResource):
     """This class represents the writing format of IndustrialCanvas.
     It is used to when data is written to CDF.
     Args:
@@ -887,7 +905,7 @@ class IndustrialCanvasApply:
                 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),
@@ -936,8 +954,52 @@ class IndustrialCanvasApply:
                 new_item_list.append(new_item)
         return new_container
 
+    @classmethod
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        """Load an IndustrialCanvasApply instance from a resource dictionary."""
+        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),
+        )
 
-class IndustrialCanvas:
+    @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]):
     """This class represents one instance of the Canvas with all connected data."""
 
     def __init__(
@@ -957,7 +1019,7 @@ class IndustrialCanvas:
         self.solution_tags = solution_tags or NodeListWithCursor[CogniteSolutionTag]([], None)
 
     @classmethod
-    def load(cls, resource: Mapping[str, list]) -> "IndustrialCanvas":
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> "IndustrialCanvas":
         """Load an IndustrialCanvas instance from a QueryResult."""
         if not ("canvas" in resource and isinstance(resource["canvas"], Sequence) and len(resource["canvas"]) == 1):
             raise ValueError("Resource does not contain a canvas node.")
@@ -1003,8 +1065,10 @@ class IndustrialCanvas:
             )
         raise TypeError(f"Expected a sequence of {node_cls.__name__}, got {type(items).__name__}")
 
-    def dump(self) -> dict[str, list]:
+    def dump(self, camel_case: bool = True) -> dict[str, list]:
         """Dump the IndustrialCanvas to a dictionary."""
+        if not camel_case:
+            raise NotImplementedError("snake_case dump is not supported")
         return {
             "canvas": [self.canvas.dump()],
             "annotations": [annotation.dump() for annotation in self.annotations],
@@ -1027,3 +1091,17 @@ class IndustrialCanvas:
             ],
             solution_tags=[solution_tag.as_write() for solution_tag in self.solution_tags],
         )
+
+    def as_id(self) -> str:
+        return self.canvas.external_id
+
+
+class IndustrialCanvasApplyList(CogniteResourceList):
+    _RESOURCE = IndustrialCanvasApply
+
+
+class IndustrialCanvasList(WriteableCogniteResourceList[IndustrialCanvasApply, IndustrialCanvas]):
+    _RESOURCE = IndustrialCanvas
+
+    def as_write(self) -> CogniteResourceList[IndustrialCanvasApply]:
+        return IndustrialCanvasApplyList([item.as_write() for item in self])

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"],
         )
 

cognite_toolkit/_cdf_tk/client/data_classes/charts_data.py

@@ -1,70 +1,30 @@
-import sys
-from dataclasses import dataclass, field, fields
-from functools import lru_cache
 from typing import Any
 
-from cognite.client import CogniteClient
-from cognite.client.data_classes._base import CogniteObject
 from cognite.client.data_classes.data_modeling import NodeId, ViewId
-from cognite.client.utils._auxiliary import to_camel_case
+from pydantic import JsonValue, field_serializer, field_validator
 
-if sys.version_info >= (3, 11):
-    from typing import Self
-else:
-    from typing_extensions import Self
+from .base import BaseModelObject
 
 
-@dataclass
-class ChartObject(CogniteObject):
-    # ChartObjects are used in the frontend and the backend does not do any validation of these fields.
-    # Therefore, to ensure that we do not lose any data, we store unknown fields in a separate dictionary.
-    # This allows unknown fields to be preserved when loading and dumping ChartObjects
-    # (serialization and deserialization).
-    _unknown_fields: dict[str, object] | None = field(default=None, init=False, repr=False)
-
-    @classmethod
-    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
-        """Load a ChartObject from a dictionary."""
-        instance = super()._load(resource, cognite_client=cognite_client)
-        instance._unknown_fields = {k: v for k, v in resource.items() if k not in cls._known_camel_case_props()}
-        return instance
-
-    @classmethod
-    @lru_cache(maxsize=1)
-    def _known_camel_case_props(cls) -> set[str]:
-        return {to_camel_case(f.name) for f in fields(cls)}
-
-    def dump(self, camel_case: bool = True) -> dict[str, Any]:
-        """Dump the ChartObject to a dictionary."""
-        data = super().dump(camel_case=camel_case)
-        if self._unknown_fields:
-            data.update(self._unknown_fields)
-        return data
-
-
-@dataclass
-class UserInfo(ChartObject):
+class UserInfo(BaseModelObject):
     id: str | None = None
     email: str | None = None
     display_name: str | None = None
 
 
-@dataclass
-class ChartSettings(ChartObject):
+class ChartSettings(BaseModelObject):
     show_y_axis: bool = True
     show_min_max: bool = True
     show_gridlines: bool = True
     merge_units: bool = False
 
 
-@dataclass
-class ThresholdFilter(ChartObject):
+class ThresholdFilter(BaseModelObject):
     min_unit: str | None = None
     max_unit: str | None = None
 
 
-@dataclass
-class ChartCall(ChartObject):
+class ChartCall(BaseModelObject):
     id: str | None = None
     hash: int | None = None
     call_id: str | None = None
@@ -72,162 +32,143 @@ class ChartCall(ChartObject):
     status: str | None = None
 
 
-@dataclass
-class SubSetting(ChartObject):
+class SubSetting(BaseModelObject):
     auto_align: bool | None = None
 
 
-@dataclass
-class FlowElement(ChartObject):
+class ChartPosition(BaseModelObject):
+    x: float | None = None
+    y: float | None = None
+
+
+class FlowElement(BaseModelObject):
     id: str | None = None
     type: str | None = None
-    position: tuple[float | None, float | None] | None = None
-    data: dict[str, object] | None = None
+    position: ChartPosition | None = None
+    data: JsonValue | None = None
+    source: str | None = None
+    target: str | None = None
+    source_handle: str | None = None
+    target_handle: str | None = None
 
 
-@dataclass
-class Flow(ChartObject):
+class Flow(BaseModelObject):
     zoom: float | None = None
     elements: list[FlowElement] | None = None
     position: tuple[float | None, float | None] | None = None
 
-    def dump(self, camel_case: bool = True) -> dict[str, Any]:
-        data = super().dump(camel_case=camel_case)
-        if self.elements:
-            data["elements"] = [el.dump(camel_case=camel_case) for el in self.elements]
-        return data
 
-    @classmethod
-    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
-        """Load a Flow object from a dictionary."""
-        instance = super()._load(resource, cognite_client=cognite_client)
-        if "elements" in resource:
-            instance.elements = [FlowElement._load(el, cognite_client=cognite_client) for el in resource["elements"]]
-        return instance
+class ChartElement(BaseModelObject):
+    id: str | None = None
+    type: str | None = None
 
 
-@dataclass
-class ChartSource(ChartObject):
-    type: str | None = None
-    id: str | None = None
+class ChartSource(ChartElement): ...
 
 
-@dataclass
-class BaseChartElement(ChartObject):
-    type: str | None = None
-    id: str | None = None
-    name: str | None = None
+class ChartCoreTimeseries(ChartElement):
+    node_reference: NodeId | None = None
+    view_reference: ViewId | None = None
+    display_mode: str | None = None
     color: str | None = None
+    created_at: int | None = None
     enabled: bool | None = None
-    line_weight: float | None = None
-    line_style: str | None = None
     interpolation: str | None = None
-    unit: str | None = None
+    line_style: str | None = None
+    line_weight: float | None = None
+    name: str | None = None
     preferred_unit: str | None = None
-    created_at: int | None = None
-    range: tuple[float | None, float | None] | None = None
-    description: str | None = None
+    range: list[float | None] | None = None
 
+    @field_serializer("node_reference", when_used="always")
+    def serialize_node_reference(self, node_reference: NodeId | None) -> dict[str, Any] | None:
+        if node_reference:
+            return node_reference.dump(include_instance_type=False)
+        return None
 
-@dataclass
-class ChartCoreTimeseries(BaseChartElement):
-    node_reference: NodeId | None = None
-    view_reference: ViewId | None = None
-    display_mode: str | None = None
+    @field_serializer("view_reference", when_used="always")
+    def serialize_view_reference(self, view_reference: ViewId | None) -> dict[str, Any] | None:
+        if view_reference:
+            return view_reference.dump(include_type=False)
+        return None
+
+    @field_validator("node_reference", mode="before")
+    @classmethod
+    def validate_node_reference(cls, value: Any) -> NodeId | None:
+        if value is None or isinstance(value, NodeId):
+            return value
+        return NodeId.load(value)
+
+    @field_validator("view_reference", mode="before")
+    @classmethod
+    def validate_view_reference(cls, value: Any) -> ViewId | None:
+        if value is None or isinstance(value, ViewId):
+            return value
+        return ViewId.load(value)
 
 
-@dataclass
-class ChartTimeseries(BaseChartElement):
+class ChartTimeseries(ChartElement):
+    color: str | None = None
+    created_at: int | None = None
+    enabled: bool | None = None
+    interpolation: str | None = None
+    line_style: str | None = None
+    line_weight: float | None = None
+    name: str | None = None
+    preferred_unit: str | None = None
+    range: list[float | None] | None = None
+    unit: str | None = None
     ts_id: int | None = None
     ts_external_id: str | None = None
     display_mode: str | None = None
     original_unit: str | None = None
+    description: str | None = None
 
 
-@dataclass
-class ChartWorkflow(BaseChartElement):
+class ChartWorkflow(ChartElement):
     version: str | None = None
+    name: str | None = None
+    color: str | None = None
+    enabled: bool | None = None
+    line_weight: float | None = None
+    line_style: str | None = None
+    interpolation: str | None = None
+    unit: str | None = None
+    preferred_unit: str | None = None
+    range: list[float | None] | None = None
+    created_at: int | None = None
     settings: SubSetting | None = None
     flow: Flow | None = None
     calls: list[ChartCall] | None = None
 
-    def dump(self, camel_case: bool = True) -> dict[str, Any]:
-        data = super().dump(camel_case=camel_case)
-        if self.settings:
-            data["settings"] = self.settings.dump(camel_case=camel_case)
-        if self.flow:
-            data["flow"] = self.flow.dump(camel_case=camel_case)
-        if self.calls:
-            data["calls"] = [c.dump(camel_case=camel_case) for c in self.calls]
-        return data
 
-    @classmethod
-    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
-        """Load a ChartWorkflow object from a dictionary."""
-        instance = super()._load(resource, cognite_client=cognite_client)
-        if "settings" in resource:
-            instance.settings = SubSetting._load(resource["settings"], cognite_client=cognite_client)
-        if "flow" in resource:
-            instance.flow = Flow._load(resource["flow"], cognite_client=cognite_client)
-        if "calls" in resource:
-            instance.calls = [ChartCall._load(call, cognite_client=cognite_client) for call in resource["calls"]]
-        return instance
-
-
-@dataclass
-class ChartThreshold(BaseChartElement):
+class ChartThreshold(ChartElement):
     visible: bool | None = None
+    name: str | None = None
     source_id: str | None = None
     upper_limit: float | None = None
     filter: ThresholdFilter | None = None
     calls: list[ChartCall] | None = None
 
-    def dump(self, camel_case: bool = True) -> dict[str, Any]:
-        data = super().dump(camel_case=camel_case)
-        if self.filter:
-            data["filter"] = self.filter.dump(camel_case=camel_case)
-        if self.calls:
-            data["calls"] = [c.dump(camel_case=camel_case) for c in self.calls]
-        return data
 
-    @classmethod
-    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
-        """Load a ChartThreshold object from a dictionary."""
-        instance = super()._load(resource, cognite_client=cognite_client)
-        if "filter" in resource:
-            instance.filter = ThresholdFilter._load(resource["filter"], cognite_client=cognite_client)
-        if "calls" in resource:
-            instance.calls = [ChartCall._load(call, cognite_client=cognite_client) for call in resource["calls"]]
-        return instance
-
-
-@dataclass
-class ChartScheduledCalculation(BaseChartElement):
+class ChartScheduledCalculation(ChartElement):
+    color: str | None = None
+    created_at: int | None = None
+    description: str | None = None
+    enabled: bool | None = None
+    interpolation: str | None = None
+    line_style: str | None = None
+    line_weight: float | None = None
+    name: str | None = None
+    preferred_unit: str | None = None
+    range: list[float | None] | None = None
+    unit: str | None = None
     version: str | None = None
     settings: SubSetting | None = None
     flow: Flow | None = None
 
-    def dump(self, camel_case: bool = True) -> dict[str, Any]:
-        data = super().dump(camel_case=camel_case)
-        if self.settings:
-            data["settings"] = self.settings.dump(camel_case=camel_case)
-        if self.flow:
-            data["flow"] = self.flow.dump(camel_case=camel_case)
-        return data
 
-    @classmethod
-    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
-        """Load a ChartScheduledCalculation object from a dictionary."""
-        instance = super()._load(resource, cognite_client=cognite_client)
-        if "settings" in resource:
-            instance.settings = SubSetting._load(resource["settings"], cognite_client=cognite_client)
-        if "flow" in resource:
-            instance.flow = Flow._load(resource["flow"], cognite_client=cognite_client)
-        return instance
-
-
-@dataclass
-class ChartData(ChartObject):
+class ChartData(BaseModelObject):
     version: int | None = None
     name: str | None = None
     date_from: str | None = None
@@ -241,62 +182,3 @@ class ChartData(ChartObject):
     threshold_collection: list[ChartThreshold] | None = None
     scheduled_calculation_collection: list[ChartScheduledCalculation] | None = None
     settings: ChartSettings | None = None
-
-    def dump(self, camel_case: bool = True) -> dict[str, Any]:
-        """Dump the ChartData object to a dictionary."""
-        data = super().dump(camel_case=camel_case)
-        list_attrs = [
-            "time_series_collection",
-            "core_timeseries_collection",
-            "workflow_collection",
-            "source_collection",
-            "threshold_collection",
-            "scheduled_calculation_collection",
-        ]
-        for attr_name in list_attrs:
-            if collection := getattr(self, attr_name):
-                key = to_camel_case(attr_name) if camel_case else attr_name
-                data[key] = [item.dump(camel_case=camel_case) for item in collection]
-
-        single_attrs_map = {
-            "user_info": "userInfo",
-            "settings": "settings",
-        }
-        for attr_name, camel_key in single_attrs_map.items():
-            if item := getattr(self, attr_name):
-                key = camel_key if camel_case else attr_name
-                data[key] = item.dump(camel_case=camel_case)
-        return data
-
-    @classmethod
-    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
-        """Load a ChartData object from a dictionary."""
-        instance = super()._load(resource, cognite_client=cognite_client)
-        collections_map = [
-            ("timeSeriesCollection", "time_series_collection", ChartTimeseries),
-            ("coreTimeseriesCollection", "core_timeseries_collection", ChartCoreTimeseries),
-            ("workflowCollection", "workflow_collection", ChartWorkflow),
-            ("sourceCollection", "source_collection", ChartSource),
-            ("thresholdCollection", "threshold_collection", ChartThreshold),
-            ("scheduledCalculationCollection", "scheduled_calculation_collection", ChartScheduledCalculation),
-        ]
-        for resource_key, attr_name, subclass in collections_map:
-            if resource_key in resource:
-                setattr(
-                    instance,
-                    attr_name,
-                    [subclass._load(item, cognite_client=cognite_client) for item in resource[resource_key]],  # type: ignore[attr-defined]
-                )
-        attribute_map = [
-            ("userInfo", "user_info", UserInfo),
-            ("settings", "settings", ChartSettings),
-        ]
-        for resource_key, attr_name, subclass in attribute_map:
-            if resource_key in resource:
-                setattr(
-                    instance,
-                    attr_name,
-                    subclass._load(resource[resource_key], cognite_client=cognite_client),  # type: ignore[attr-defined]
-                )
-
-        return instance