Add test to ensure CogniteObject subclasses handle unknown args when deserializing (#1504)

Author: erlendvollset

.pre-commit-config.yaml

@@ -1,9 +1,9 @@
 ---
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.1.5
     hooks:
       - id: ruff
-        rev: v0.1.5
         args:
           - --fix
           - --exit-non-zero-on-fix

cognite/client/_api/geospatial.py

@@ -943,10 +943,12 @@ def put_raster(
             self._raster_resource_path(feature_type_external_id, feature_external_id, raster_property_name)
             + f"?{query_params}"
         )
+        with open(file, "rb") as fh:
+            data = fh.read()
         res = self._do_request(
             "PUT",
             url_path,
-            data=open(file, "rb").read(),
+            data=data,
             headers={"Content-Type": "application/binary"},
             timeout=self._config.timeout,
         )

cognite/client/data_classes/_base.py

@@ -119,9 +119,6 @@ class CogniteObject:
     It is used both by the CogniteResources and the nested classes used by the CogniteResources.
     """
 
-    def __init__(self, cognite_client: CogniteClient | None = None) -> None:
-        raise NotImplementedError
-
     def __eq__(self, other: Any) -> bool:
         return type(self) is type(other) and self.dump() == other.dump()
 

cognite/client/data_classes/capabilities.py

@@ -1167,7 +1167,7 @@ class Scope:
     type[Capability], tuple[tuple[type[Capability.Scope], ...], tuple[str, ...]]
 ] = MappingProxyType(
     {
-        acl: tuple(zip(*[(v, k) for k, v in vars(acl.Scope).items() if inspect.isclass(v)]))  # type: ignore [misc]
+        acl: tuple(zip(*[(v, k) for k, v in vars(acl.Scope).items() if inspect.isclass(v)]))
         for acl in _CAPABILITY_CLASS_BY_NAME.values()
     }
 )

cognite/client/data_classes/data_modeling/containers.py

@@ -2,16 +2,17 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import asdict, dataclass
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 from typing_extensions import Self
 
 from cognite.client.data_classes._base import (
     CogniteFilter,
+    CogniteObject,
     CogniteResourceList,
 )
-from cognite.client.data_classes.data_modeling._core import DataModelingResource
 from cognite.client.data_classes.data_modeling._validation import validate_data_modeling_identifier
+from cognite.client.data_classes.data_modeling.core import DataModelingResource
 from cognite.client.data_classes.data_modeling.data_types import (
     DirectRelation,
     PropertyType,
@@ -112,6 +113,23 @@ def __init__(
         validate_data_modeling_identifier(space, external_id)
         super().__init__(space, external_id, properties, description, name, used_for, constraints, indexes)
 
+    @classmethod
+    def _load(cls, resource: dict, cognite_client: CogniteClient | None = None) -> ContainerApply:
+        return ContainerApply(
+            space=resource["space"],
+            external_id=resource["externalId"],
+            properties={k: ContainerProperty.load(v) for k, v in resource["properties"].items()},
+            description=resource.get("description"),
+            name=resource.get("name"),
+            used_for=resource.get("usedFor"),
+            constraints={k: Constraint.load(v) for k, v in resource["constraints"].items()}
+            if "constraints" in resource
+            else None,
+            indexes={k: Index.load(v) for k, v in resource["indexes"].items()} or None
+            if "indexes" in resource
+            else None,
+        )
+
 
 class Container(ContainerCore):
     """Represent the physical storage of data. This is the read format of the container
@@ -210,7 +228,7 @@ def __init__(self, space: str | None = None, include_global: bool = False) -> No
 
 
 @dataclass(frozen=True)
-class ContainerProperty:
+class ContainerProperty(CogniteObject):
     type: PropertyType
     nullable: bool = True
     auto_increment: bool = False
@@ -219,21 +237,21 @@ class ContainerProperty:
     description: str | None = None
 
     @classmethod
-    def load(cls, data: dict[str, Any]) -> ContainerProperty:
-        if "type" not in data:
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        if "type" not in resource:
             raise ValueError("Type not specified")
-        if data["type"].get("type") == "direct":
-            type_: PropertyType = DirectRelation.load(data["type"])
+        if resource["type"].get("type") == "direct":
+            type_: PropertyType = DirectRelation.load(resource["type"])
         else:
-            type_ = PropertyType.load(data["type"])
+            type_ = PropertyType.load(resource["type"])
         return cls(
             type=type_,
             # If nothing is specified, we will pass through null values
-            nullable=data.get("nullable"),  # type: ignore[arg-type]
-            auto_increment=data.get("autoIncrement"),  # type: ignore[arg-type]
-            name=data.get("name"),
-            default_value=data.get("defaultValue"),
-            description=data.get("description"),
+            nullable=resource.get("nullable"),  # type: ignore[arg-type]
+            auto_increment=resource.get("autoIncrement"),  # type: ignore[arg-type]
+            name=resource.get("name"),
+            default_value=resource.get("defaultValue"),
+            description=resource.get("description"),
         )
 
     def dump(self, camel_case: bool = True) -> dict[str, str | dict]:
@@ -247,14 +265,14 @@ def dump(self, camel_case: bool = True) -> dict[str, str | dict]:
 
 
 @dataclass(frozen=True)
-class Constraint(ABC):
+class Constraint(CogniteObject, ABC):
     @classmethod
-    def load(cls, data: dict) -> RequiresConstraint | UniquenessConstraint:
-        if data["constraintType"] == "requires":
-            return RequiresConstraint.load(data)
-        elif data["constraintType"] == "uniqueness":
-            return UniquenessConstraint.load(data)
-        raise ValueError(f"Invalid constraint type {data['constraintType']}")
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        if resource["constraintType"] == "requires":
+            return cast(Self, RequiresConstraint.load(resource))
+        elif resource["constraintType"] == "uniqueness":
+            return cast(Self, UniquenessConstraint.load(resource))
+        raise ValueError(f"Invalid constraint type {resource['constraintType']}")
 
     @abstractmethod
     def dump(self, camel_case: bool = True) -> dict[str, str | dict]:
@@ -266,8 +284,8 @@ class RequiresConstraint(Constraint):
     require: ContainerId
 
     @classmethod
-    def load(cls, data: dict) -> RequiresConstraint:
-        return cls(require=ContainerId.load(data["require"]))
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        return cls(require=ContainerId.load(resource["require"]))
 
     def dump(self, camel_case: bool = True) -> dict[str, str | dict]:
         as_dict = asdict(self)
@@ -284,8 +302,8 @@ class UniquenessConstraint(Constraint):
     properties: list[str]
 
     @classmethod
-    def load(cls, data: dict) -> UniquenessConstraint:
-        return cls(properties=data["properties"])
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        return cls(properties=resource["properties"])
 
     def dump(self, camel_case: bool = True) -> dict[str, str | dict]:
         as_dict = asdict(self)
@@ -296,14 +314,14 @@ def dump(self, camel_case: bool = True) -> dict[str, str | dict]:
 
 
 @dataclass(frozen=True)
-class Index(ABC):
+class Index(CogniteObject, ABC):
     @classmethod
-    def load(cls, data: dict) -> Index:
-        if data["indexType"] == "btree":
-            return BTreeIndex.load(data)
-        if data["indexType"] == "inverted":
-            return InvertedIndex.load(data)
-        raise ValueError(f"Invalid index type {data['indexType']}")
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        if resource["indexType"] == "btree":
+            return cast(Self, BTreeIndex.load(resource))
+        if resource["indexType"] == "inverted":
+            return cast(Self, InvertedIndex.load(resource))
+        raise ValueError(f"Invalid index type {resource['indexType']}")
 
     @abstractmethod
     def dump(self, camel_case: bool = True) -> dict[str, str | dict]:
@@ -316,8 +334,8 @@ class BTreeIndex(Index):
     cursorable: bool = False
 
     @classmethod
-    def load(cls, data: dict[str, Any]) -> BTreeIndex:
-        return cls(properties=data["properties"], cursorable=data.get("cursorable"))  # type: ignore[arg-type]
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        return cls(properties=resource["properties"], cursorable=resource.get("cursorable"))  # type: ignore[arg-type]
 
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
         dumped: dict[str, Any] = {"properties": self.properties}
@@ -332,8 +350,8 @@ class InvertedIndex(Index):
     properties: list[str]
 
     @classmethod
-    def load(cls, data: dict[str, Any]) -> InvertedIndex:
-        return cls(properties=data["properties"])
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        return cls(properties=resource["properties"])
 
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
         dumped: dict[str, Any] = {"properties": self.properties}

cognite/client/data_classes/data_modeling/_core.py renamed to cognite/client/data_classes/data_modeling/core.py

@@ -6,7 +6,7 @@
 
 from typing_extensions import Self
 
-from cognite.client.data_classes._base import CogniteResource, basic_instance_dump
+from cognite.client.data_classes._base import CogniteObject, CogniteResource, basic_instance_dump
 from cognite.client.utils._auxiliary import json_dump_default
 from cognite.client.utils._text import convert_all_keys_to_snake_case
 
@@ -34,13 +34,14 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
         return cls(**convert_all_keys_to_snake_case(resource))
 
 
-class DataModelingSort:
+class DataModelingSort(CogniteObject):
     def __init__(
         self,
         property: str | list[str] | tuple[str, ...],
         direction: Literal["ascending", "descending"] = "ascending",
         nulls_first: bool = False,
     ) -> None:
+        super().__init__()
         self.property = property
         self.direction = direction
         self.nulls_first = nulls_first
@@ -55,7 +56,7 @@ def __repr__(self) -> str:
         return str(self)
 
     @classmethod
-    def load(cls, resource: dict) -> Self:
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
         if not isinstance(resource, dict):
             raise TypeError(f"Resource must be mapping, not {type(resource)}")
 

cognite/client/data_classes/data_modeling/data_models.py

@@ -6,8 +6,8 @@
 from typing_extensions import Self
 
 from cognite.client.data_classes._base import CogniteFilter, CogniteResourceList
-from cognite.client.data_classes.data_modeling._core import DataModelingResource, DataModelingSort
 from cognite.client.data_classes.data_modeling._validation import validate_data_modeling_identifier
+from cognite.client.data_classes.data_modeling.core import DataModelingResource, DataModelingSort
 from cognite.client.data_classes.data_modeling.ids import DataModelId, ViewId
 from cognite.client.data_classes.data_modeling.views import View, ViewApply
 
@@ -79,10 +79,15 @@ def _load_view(cls, view_data: dict) -> ViewId | ViewApply:
             return ViewApply._load(view_data)
 
     @classmethod
-    def _load(cls, resource: dict, cognite_client: CogniteClient | None = None) -> Self:
-        if "views" in resource:
-            resource["views"] = [cls._load_view(v) for v in resource["views"]] or None
-        return super()._load(resource)
+    def _load(cls, resource: dict, cognite_client: CogniteClient | None = None) -> DataModelApply:
+        return DataModelApply(
+            space=resource["space"],
+            external_id=resource["externalId"],
+            version=resource["version"],
+            description=resource.get("description"),
+            name=resource.get("name"),
+            views=[cls._load_view(v) for v in resource["views"]] if "views" in resource else None,
+        )
 
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
         output = super().dump(camel_case)

cognite/client/data_classes/data_modeling/graphql.py

@@ -1,13 +1,19 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
+from typing_extensions import Self
+
+from cognite.client.data_classes._base import CogniteObject
 from cognite.client.data_classes.data_modeling.ids import DataModelId
 
+if TYPE_CHECKING:
+    from cognite.client import CogniteClient
+
 
 @dataclass
-class DMLApplyResult:
+class DMLApplyResult(CogniteObject):
     space: str
     external_id: str
     version: str
@@ -24,13 +30,13 @@ def as_id(self) -> DataModelId:
         )
 
     @classmethod
-    def load(cls, data: dict[str, Any]) -> DMLApplyResult:
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
         return cls(
-            space=data["space"],
-            external_id=data["externalId"],
-            version=data["version"],
-            name=data["name"],
-            description=data["description"],
-            created_time=data["createdTime"],
-            last_updated_time=data["lastUpdatedTime"],
+            space=resource["space"],
+            external_id=resource["externalId"],
+            version=resource["version"],
+            name=resource["name"],
+            description=resource["description"],
+            created_time=resource["createdTime"],
+            last_updated_time=resource["lastUpdatedTime"],
         )