Support node.type base property on nodes (#1444)

cognitedata · Oct 27, 2023 · 1b20f3d · 1b20f3d
1 parent 45c502d
commit 1b20f3d
Show file tree

Hide file tree

Showing 5 changed files with 122 additions and 45 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -17,6 +17,10 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## [6.37.0] - 2023-10-27
+### Added
+- Support for `type` property in `NodeApply` and `Node`.
+
 ## [6.36.0] - 2023-10-25
 ### Added
 - Support for listing members of Data Point Subscription, `client.time_series.subscriptions.list_member_time_series()`. Note this is an experimental feature.

diff --git a/cognite/client/_version.py b/cognite/client/_version.py
@@ -1,4 +1,4 @@
 from __future__ import annotations
 
-__version__ = "6.36.0"
+__version__ = "6.37.0"
 __api_subversion__ = "V20220125"
diff --git a/cognite/client/data_classes/data_modeling/instances.py b/cognite/client/data_classes/data_modeling/instances.py
@@ -26,6 +26,8 @@
     overload,
 )
 
+from typing_extensions import Self
+
 from cognite.client.data_classes._base import CogniteResourceList
 from cognite.client.data_classes.aggregations import AggregatedNumberedValue
 from cognite.client.data_classes.data_modeling._core import DataModelingResource, DataModelingSort
@@ -108,7 +110,7 @@ class InstanceCore(DataModelingResource):
         instance_type (Literal["node", "edge"]): No description.
     """
 
-    def __init__(self, space: str, external_id: str, instance_type: Literal["node", "edge"] = "node") -> None:
+    def __init__(self, space: str, external_id: str, instance_type: Literal["node", "edge"]) -> None:
         self.instance_type = instance_type
         self.space = space
         self.external_id = external_id
@@ -145,7 +147,7 @@ def dump(self, camel_case: bool = False) -> dict[str, Any]:
         return output
 
     @classmethod
-    def load(cls: type[T_Instance_Apply], data: dict | str) -> T_Instance_Apply:
+    def load(cls, data: dict | str) -> Self:
         data = data if isinstance(data, dict) else json.loads(data)
         data = convert_all_keys_to_snake_case(data)
         if cls is not InstanceApply:
@@ -157,8 +159,6 @@ def load(cls: type[T_Instance_Apply], data: dict | str) -> T_Instance_Apply:
         return instance
 
 
-T_Instance_Apply = TypeVar("T_Instance_Apply", bound=InstanceApply)
-
 _T = TypeVar("_T")
 
 
@@ -185,7 +185,7 @@ def dump(self) -> dict[Space, dict[str, dict[PropertyIdentifier, PropertyValue]]
         for view_id, properties in self.data.items():
             view_id_str = f"{view_id.external_id}/{view_id.version}"
             props[view_id.space][view_id_str] = cast(Dict[PropertyIdentifier, PropertyValue], properties)
-        return props
+        return dict(props)
 
     def items(self) -> ItemsView[ViewId, MutableMapping[PropertyIdentifier, PropertyValue]]:
         return self.data.items()
@@ -237,16 +237,13 @@ def __setitem__(self, view: ViewIdentifier, properties: MutableMapping[PropertyI
         self.data[view_id] = properties
 
 
-T_Instance = TypeVar("T_Instance", bound="Instance")
-
-
 class Instance(InstanceCore):
     """A node or edge. This is the read version of the instance.
 
     Args:
         space (str): The workspace for the instance, a unique identifier for the space.
         external_id (str): Combined with the space is the unique identifier of the instance.
-        version (str): DMS version.
+        version (int): DMS version.
         last_updated_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
         created_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
         instance_type (Literal["node", "edge"]): The type of instance.
@@ -259,7 +256,7 @@ def __init__(
         self,
         space: str,
         external_id: str,
-        version: str,
+        version: int,
         last_updated_time: int,
         created_time: int,
         instance_type: Literal["node", "edge"] = "node",
@@ -275,7 +272,7 @@ def __init__(
         self.properties: Properties = properties or Properties({})
 
     @classmethod
-    def load(cls: type[T_Instance], data: dict | str) -> T_Instance:
+    def load(cls, data: dict | str) -> Self:
         data = json.loads(data) if isinstance(data, str) else data
         if "properties" in data:
             data["properties"] = Properties.load(data["properties"])
@@ -301,7 +298,7 @@ class InstanceApplyResult(InstanceCore):
         instance_type (Literal["node", "edge"]): The type of instance.
         space (str): The workspace for the instance, a unique identifier for the space.
         external_id (str): Combined with the space is the unique identifier of the instance.
-        version (str): DMS version of the instance.
+        version (int): DMS version of the instance.
         was_modified (bool): Whether the instance was modified by the ingestion.
         last_updated_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
         created_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
@@ -313,7 +310,7 @@ def __init__(
         instance_type: Literal["node", "edge"],
         space: str,
         external_id: str,
-        version: str,
+        version: int,
         was_modified: bool,
         last_updated_time: int,
         created_time: int,
@@ -386,6 +383,7 @@ class NodeApply(InstanceApply):
         external_id (str): Combined with the space is the unique identifier of the node.
         existing_version (int | None): Fail the ingestion request if the node's version is greater than or equal to this value. If no existingVersion is specified, the ingestion will always overwrite any existing data for the edge (for the specified container or node). If existingVersion is set to 0, the upsert will behave as an insert, so it will fail the bulk if the item already exists. If skipOnVersionConflict is set on the ingestion request, then the item will be skipped instead of failing the ingestion request.
         sources (list[NodeOrEdgeData] | None): List of source properties to write. The properties are from the node and/or container the container(s) making up this node.
+        type (DirectRelationReference | tuple[str, str] | None): Direct relation pointing to the type node.
     """
 
     def __init__(
@@ -394,8 +392,26 @@ def __init__(
         external_id: str,
         existing_version: int | None = None,
         sources: list[NodeOrEdgeData] | None = None,
+        type: DirectRelationReference | tuple[str, str] | None = None,
     ) -> None:
         super().__init__(space, external_id, "node", existing_version, sources)
+        if isinstance(type, tuple):
+            self.type: DirectRelationReference | None = DirectRelationReference.load(type)
+        else:
+            self.type = type
+
+    def dump(self, camel_case: bool = False) -> dict[str, Any]:
+        output = super().dump(camel_case)
+        if self.type:
+            output["type"] = self.type.dump(camel_case)
+        return output
+
+    @classmethod
+    def load(cls, data: dict | str) -> NodeApply:
+        data = json.loads(data) if isinstance(data, str) else data
+        instance = super().load(data)
+        instance.type = DirectRelationReference.load(data["type"]) if "type" in data else None
+        return instance
 
     def as_id(self) -> NodeId:
         return NodeId(space=self.space, external_id=self.external_id)
@@ -407,26 +423,29 @@ class Node(Instance):
     Args:
         space (str): The workspace for the node, a unique identifier for the space.
         external_id (str): Combined with the space is the unique identifier of the node.
-        version (str): DMS version.
+        version (int): DMS version.
         last_updated_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
         created_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
         deleted_time (int | None): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds. Timestamp when the instance was soft deleted. Note that deleted instances are filtered out of query results, but present in sync results
         properties (Properties | None): Properties of the node.
+        type (DirectRelationReference | None): Direct relation pointing to the type node.
         **_ (Any): No description.
     """
 
     def __init__(
         self,
         space: str,
         external_id: str,
-        version: str,
+        version: int,
         last_updated_time: int,
         created_time: int,
-        deleted_time: int | None = None,
-        properties: Properties | None = None,
+        deleted_time: int | None,
+        properties: Properties | None,
+        type: DirectRelationReference | None,
         **_: Any,
     ) -> None:
         super().__init__(space, external_id, version, last_updated_time, created_time, "node", deleted_time, properties)
+        self.type = type
 
     def as_apply(self, source: ViewIdentifier | ContainerIdentifier, existing_version: int) -> NodeApply:
         """
@@ -457,14 +476,34 @@ def as_apply(self, source: ViewIdentifier | ContainerIdentifier, existing_versio
     def as_id(self) -> NodeId:
         return NodeId(space=self.space, external_id=self.external_id)
 
+    def dump(self, camel_case: bool = False) -> dict[str, Any]:
+        output = super().dump(camel_case)
+        if self.type:
+            output["type"] = self.type.dump(camel_case)
+        return output
+
+    @classmethod
+    def load(cls, data: dict | str) -> Node:
+        data = json.loads(data) if isinstance(data, str) else data
+        return Node(
+            space=data["space"],
+            external_id=data["externalId"],
+            version=data["version"],
+            last_updated_time=data["lastUpdatedTime"],
+            created_time=data["createdTime"],
+            deleted_time=data.get("deletedTime"),
+            properties=Properties.load(data["properties"]) if "properties" in data else None,
+            type=DirectRelationReference.load(data["type"]) if "type" in data else None,
+        )
+
 
 class NodeApplyResult(InstanceApplyResult):
     """A node. This represents the update on the node.
 
     Args:
         space (str): The workspace for the node, a unique identifier for the space.
         external_id (str): Combined with the space is the unique identifier of the node.
-        version (str): DMS version of the node.
+        version (int): DMS version of the node.
         was_modified (bool): Whether the node was modified by the ingestion.
         last_updated_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
         created_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
@@ -475,7 +514,7 @@ def __init__(
         self,
         space: str,
         external_id: str,
-        version: str,
+        version: int,
         was_modified: bool,
         last_updated_time: int,
         created_time: int,
@@ -557,7 +596,7 @@ class Edge(Instance):
     Args:
         space (str): The workspace for the edge, a unique identifier for the space.
         external_id (str): Combined with the space is the unique identifier of the edge.
-        version (str): DMS version.
+        version (int): DMS version.
         type (DirectRelationReference): The type of edge.
         last_updated_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
         created_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
@@ -572,7 +611,7 @@ def __init__(
         self,
         space: str,
         external_id: str,
-        version: str,
+        version: int,
         type: DirectRelationReference,
         last_updated_time: int,
         created_time: int,
@@ -644,7 +683,7 @@ class EdgeApplyResult(InstanceApplyResult):
     Args:
         space (str): The workspace for the edge, a unique identifier for the space.
         external_id (str): Combined with the space is the unique identifier of the edge.
-        version (str): DMS version.
+        version (int): DMS version.
         was_modified (bool): Whether the edge was modified by the ingestion.
         last_updated_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
         created_time (int): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
@@ -655,7 +694,7 @@ def __init__(
         self,
         space: str,
         external_id: str,
-        version: str,
+        version: int,
         was_modified: bool,
         last_updated_time: int,
         created_time: int,

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "cognite-sdk"
 
-version = "6.36.0"
+version = "6.37.0"
 description = "Cognite Python SDK"
 readme = "README.md"
 documentation = "https://cognite-sdk-python.readthedocs-hosted.com"

diff --git a/tests/tests_unit/test_data_classes/test_data_models/test_instances.py b/tests/tests_unit/test_data_classes/test_data_models/test_instances.py
@@ -2,14 +2,15 @@
     ContainerId,
     DirectRelationReference,
     EdgeApply,
+    Node,
     NodeApply,
     NodeId,
     NodeOrEdgeData,
 )
 
 
 class TestEdgeApply:
-    def test_dump(self) -> None:
+    def test_dump_and_load(self) -> None:
         edge = EdgeApply(
             space="mySpace",
             external_id="relation:arnold_schwarzenegger:actor",
@@ -18,7 +19,7 @@ def test_dump(self) -> None:
             end_node=DirectRelationReference("mySpace", "actor.external_id"),
         )
 
-        assert edge.dump(camel_case=True) == {
+        assert EdgeApply.load(edge.dump(camel_case=True)).dump(camel_case=True) == {
             "space": "mySpace",
             "externalId": "relation:arnold_schwarzenegger:actor",
             "type": {
@@ -55,11 +56,11 @@ def test_direct_relation_serialization(self) -> None:
 
 
 class TestNodeApply:
-    def test_dump_with_snake_case_fields(self) -> None:
-        # Arrange
+    def test_dump_and_load(self) -> None:
         node = NodeApply(
             space="IntegrationTestsImmutable",
             external_id="shop:case:integration_test",
+            type=("someSpace", "someType"),
             sources=[
                 NodeOrEdgeData(
                     source=ContainerId("IntegrationTestsImmutable", "Case"),
@@ -82,21 +83,54 @@ def test_dump_with_snake_case_fields(self) -> None:
             ],
         )
 
-        # Act
-        dumped = node.dump(camel_case=True)
+        assert NodeApply.load(node.dump(camel_case=True)).dump(camel_case=True) == {
+            "externalId": "shop:case:integration_test",
+            "instanceType": "node",
+            "sources": [
+                {
+                    "properties": {
+                        "arguments": "Integration test",
+                        "bid": "shop:bid_matrix:8",
+                        "bid_history": ["shop:bid_matrix:9"],
+                        "commands": {
+                            "externalId": "shop:command_config:integration_test",
+                            "space": "IntegrationTestsImmutable",
+                        },
+                        "cut_files": ["shop:cut_file:1"],
+                        "end_time": "2021-01-01T00:00:00",
+                        "name": "Integration test",
+                        "runStatus": "Running",
+                        "scenario": "Integration test",
+                        "start_time": "2021-01-01T00:00:00",
+                    },
+                    "source": {"externalId": "Case", "space": "IntegrationTestsImmutable", "type": "container"},
+                }
+            ],
+            "space": "IntegrationTestsImmutable",
+            "type": {"externalId": "someType", "space": "someSpace"},
+        }
 
-        # Assert
-        assert sorted(dumped["sources"][0]["properties"]) == sorted(
-            [
-                "name",
-                "scenario",
-                "start_time",
-                "end_time",
-                "cut_files",
-                "bid",
-                "bid_history",
-                "runStatus",
-                "arguments",
-                "commands",
-            ]
+
+class TestNode:
+    def test_dump_and_load(self) -> None:
+        node = Node(
+            space="IntegrationTestsImmutable",
+            external_id="shop:case:integration_test",
+            version=1,
+            type=DirectRelationReference("someSpace", "someType"),
+            last_updated_time=123,
+            created_time=123,
+            deleted_time=None,
+            properties=None,
         )
+
+        assert Node.load(node.dump(camel_case=True)).dump(camel_case=True) == {
+            "createdTime": 123,
+            "externalId": "shop:case:integration_test",
+            "instanceType": "node",
+            "lastUpdatedTime": 123,
+            "properties": {},
+            "space": "IntegrationTestsImmutable",
+            "type": {"externalId": "someType", "space": "someSpace"},
+            "version": 1,
+        }