cognite-toolkit 0.7.70__py3-none-any.whl → 0.7.71__py3-none-any.whl

cognite_toolkit/_cdf_tk/client/resource_classes/infield.py

@@ -1,35 +1,38 @@
-import sys
 from typing import Any, ClassVar, Literal
 
-from pydantic import JsonValue, field_validator
-from pydantic_core.core_schema import ValidationInfo
+from pydantic import JsonValue, model_validator
 
-from cognite_toolkit._cdf_tk.client._resource_base import BaseResourceList, ResponseResource
-from cognite_toolkit._cdf_tk.protocols import (
-    ResourceRequestListProtocol,
-    ResourceResponseListProtocol,
-)
+from cognite_toolkit._cdf_tk.client._resource_base import BaseModelObject
 from cognite_toolkit._cdf_tk.utils.text import sanitize_instance_external_id
 
-from .instance_api import InstanceRequestResource, ViewReference
-
-if sys.version_info >= (3, 11):
-    from typing import Self
-else:
-    from typing_extensions import Self
+from .instance_api import (
+    TypedInstanceIdentifier,
+    TypedNodeIdentifier,
+    TypedViewReference,
+    WrappedInstanceListRequest,
+    WrappedInstanceListResponse,
+    WrappedInstanceRequest,
+    WrappedInstanceResponse,
+    move_properties,
+)
 
-INFIELD_LOCATION_CONFIG_VIEW_ID = ViewReference(space="cdf_infield", external_id="InFieldLocationConfig", version="v1")
-INFIELD_CDM_LOCATION_CONFIG_VIEW_ID = ViewReference(
+INFIELD_LOCATION_CONFIG_VIEW_ID = TypedViewReference(
+    space="cdf_infield", external_id="InFieldLocationConfig", version="v1"
+)
+INFIELD_CDM_LOCATION_CONFIG_VIEW_ID = TypedViewReference(
     space="infield_cdm_source_desc_sche_asset_file_ts", external_id="InFieldCDMLocationConfig", version="v1"
 )
-DATA_EXPLORATION_CONFIG_VIEW_ID = ViewReference(space="cdf_infield", external_id="DataExplorationConfig", version="v1")
+DATA_EXPLORATION_CONFIG_VIEW_ID = TypedViewReference(
+    space="cdf_infield", external_id="DataExplorationConfig", version="v1"
+)
 
 
-class DataExplorationConfig(InstanceRequestResource):
+class DataExplorationConfig(BaseModelObject):
     """Data Exploration Configuration resource class."""
 
-    VIEW_ID: ClassVar[ViewReference] = DATA_EXPLORATION_CONFIG_VIEW_ID
-    instance_type: Literal["node"] = "node"
+    VIEW_ID: ClassVar[TypedViewReference] = DATA_EXPLORATION_CONFIG_VIEW_ID
+    space: str | None = None
+    external_id: str | None = None
 
     observations: dict[str, JsonValue] | None = None
     activities: dict[str, JsonValue] | None = None
@@ -37,19 +40,19 @@ class DataExplorationConfig(InstanceRequestResource):
     notifications: dict[str, JsonValue] | None = None
     assets: dict[str, JsonValue] | None = None
 
+    @model_validator(mode="before")
+    @classmethod
+    def move_properties(cls, data: dict[str, Any]) -> dict[str, Any]:
+        return move_properties(data, cls.VIEW_ID)
 
-class InfieldLocationConfig(
-    ResponseResource["InfieldLocationConfig"],
-    InstanceRequestResource,
-):
+
+class InFieldLocationConfig(BaseModelObject):
     """Infield Location Configuration resource class.
 
     This class is used for both the response and request resource for Infield Location Configuration nodes.
     """
 
-    VIEW_ID: ClassVar[ViewReference] = INFIELD_LOCATION_CONFIG_VIEW_ID
-    instance_type: Literal["node"] = "node"
-
+    VIEW_ID: ClassVar[TypedViewReference] = INFIELD_LOCATION_CONFIG_VIEW_ID
     root_location_external_id: str | None = None
     feature_toggles: dict[str, JsonValue] | None = None
     app_instance_space: str | None = None
@@ -57,48 +60,94 @@ class InfieldLocationConfig(
     data_filters: dict[str, JsonValue] | None = None
     data_exploration_config: DataExplorationConfig | None = None
 
-    def as_request_resource(self) -> "InfieldLocationConfig":
-        return self
-
-    def as_write(self) -> Self:
-        return self
-
-    @field_validator("data_exploration_config", mode="before")
-    @classmethod
-    def generate_identifier_if_missing(cls, value: Any, info: ValidationInfo) -> Any:
-        """We do not require the user to specify the space and externalId for the data exploration config."""
-        if isinstance(value, dict):
-            if value.get("space") is None:
-                value["space"] = info.data["space"]
-            if value.get("externalId") is None:
-                external_id = info.data["external_id"]
-                candidate = f"{external_id}_data_exploration_config"
-                value["externalId"] = sanitize_instance_external_id(candidate)
-        return value
-
-
-class InfieldLocationConfigList(
-    BaseResourceList[InfieldLocationConfig],
-    ResourceResponseListProtocol,
-    ResourceRequestListProtocol,
-):
-    """A list of InfieldLocationConfig objects."""
-
-    _RESOURCE = InfieldLocationConfig
-
-    def as_write(self) -> Self:
-        return self
-
-
-class InFieldCDMLocationConfig(ResponseResource["InFieldCDMLocationConfig"], InstanceRequestResource):
-    """InField CDM Location Configuration resource class.
-
-    This class is used for both the response and request resource for InField CDM Location Configuration nodes.
-    """
-
-    VIEW_ID: ClassVar[ViewReference] = INFIELD_CDM_LOCATION_CONFIG_VIEW_ID
-    instance_type: Literal["node"] = "node"
 
+class InFieldLocationConfigRequest(WrappedInstanceListRequest, InFieldLocationConfig):
+    def dump_instances(self) -> list[dict[str, Any]]:
+        space: str | None = None
+        external_id: str | None = None
+        if self.data_exploration_config:
+            space = self.data_exploration_config.space or self.space
+            if self.data_exploration_config.external_id:
+                external_id = self.data_exploration_config.external_id
+            else:
+                candidate = f"{self.external_id}_data_exploration_config"
+                external_id = sanitize_instance_external_id(candidate)
+
+        properties = self.model_dump(
+            by_alias=True,
+            exclude_unset=True,
+            exclude={"data_exploration_config", "instance_type", "space", "external_id"},
+        )
+        if space and external_id:
+            properties["dataExplorationConfig"] = {"space": space, "externalId": external_id}
+        output: list[dict[str, Any]] = [
+            {
+                "instanceType": self.instance_type,
+                "space": self.space,
+                "externalId": self.external_id,
+                "sources": [
+                    {
+                        "source": self.VIEW_ID.dump(),
+                        "properties": properties,
+                    }
+                ],
+            }
+        ]
+        if space and external_id and self.data_exploration_config:
+            output.append(
+                {
+                    "instanceType": "node",
+                    "space": space,
+                    "externalId": external_id,
+                    "sources": [
+                        {
+                            "source": DataExplorationConfig.VIEW_ID.dump(),
+                            "properties": self.data_exploration_config.model_dump(
+                                by_alias=True, exclude_unset=True, exclude={"space", "external_id"}
+                            ),
+                        }
+                    ],
+                }
+            )
+        return output
+
+    def as_ids(self) -> list[TypedInstanceIdentifier]:
+        output: list[TypedInstanceIdentifier] = [self.as_id()]
+        if (
+            self.data_exploration_config
+            and self.data_exploration_config.space
+            and self.data_exploration_config.external_id
+        ):
+            output.append(
+                TypedNodeIdentifier(
+                    space=self.data_exploration_config.space,
+                    external_id=self.data_exploration_config.external_id,
+                )
+            )
+        return output
+
+
+class InFieldLocationConfigResponse(WrappedInstanceListResponse, InFieldLocationConfig):
+    def as_request_resource(self) -> InFieldLocationConfigRequest:
+        return InFieldLocationConfigRequest.model_validate(self.dump(), extra="ignore")
+
+    def as_ids(self) -> list[TypedInstanceIdentifier]:
+        output: list[TypedInstanceIdentifier] = [TypedNodeIdentifier(space=self.space, external_id=self.external_id)]
+        if (
+            self.data_exploration_config
+            and self.data_exploration_config.space
+            and self.data_exploration_config.external_id
+        ):
+            output.append(
+                TypedNodeIdentifier(
+                    space=self.data_exploration_config.space,
+                    external_id=self.data_exploration_config.external_id,
+                )
+            )
+        return output
+
+
+class InFieldCDMLocationConfig(BaseModelObject):
     name: str | None = None
     description: str | None = None
     feature_toggles: dict[str, JsonValue] | None = None
@@ -109,8 +158,20 @@ class InFieldCDMLocationConfig(ResponseResource["InFieldCDMLocationConfig"], Ins
     disciplines: list[dict[str, JsonValue]] | None = None
     data_exploration_config: dict[str, JsonValue] | None = None
 
-    def as_request_resource(self) -> "InFieldCDMLocationConfig":
-        return self
 
-    def as_write(self) -> Self:
-        return self
+class InFieldCDMLocationConfigRequest(WrappedInstanceRequest, InFieldCDMLocationConfig):
+    VIEW_ID: ClassVar[TypedViewReference] = INFIELD_CDM_LOCATION_CONFIG_VIEW_ID
+    instance_type: Literal["node"] = "node"
+
+    def as_id(self) -> TypedNodeIdentifier:
+        return TypedNodeIdentifier(space=self.space, external_id=self.external_id)
+
+
+class InFieldCDMLocationConfigResponse(
+    WrappedInstanceResponse[InFieldCDMLocationConfigRequest], InFieldCDMLocationConfig
+):
+    VIEW_ID: ClassVar[TypedViewReference] = INFIELD_CDM_LOCATION_CONFIG_VIEW_ID
+    instance_type: Literal["node"] = "node"
+
+    def as_request_resource(self) -> InFieldCDMLocationConfigRequest:
+        return InFieldCDMLocationConfigRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/resource_classes/instance_api.py

@@ -1,8 +1,14 @@
-from typing import Any, ClassVar, Literal, TypeAlias
+from abc import ABC, abstractmethod
+from typing import Any, ClassVar, Literal, TypeAlias, TypeVar
 
-from pydantic import ConfigDict, JsonValue, model_serializer
+from pydantic import model_validator
 
-from cognite_toolkit._cdf_tk.client._resource_base import BaseModelObject, Identifier, RequestResource
+from cognite_toolkit._cdf_tk.client._resource_base import (
+    BaseModelObject,
+    Identifier,
+    RequestResource,
+    ResponseResource,
+)
 
 InstanceType: TypeAlias = Literal["node", "edge"]
 
@@ -39,6 +45,9 @@ class TypedEdgeIdentifier(TypedInstanceIdentifier):
         return f"Edge({self.space}, {self.external_id})"
 
 
+T_TypedInstanceIdentifier = TypeVar("T_TypedInstanceIdentifier", bound=TypedInstanceIdentifier)
+
+
 class InstanceIdentifier(Identifier):
     space: str
     external_id: str
@@ -64,7 +73,7 @@ class InstanceResult(BaseModelObject):
         )
 
 
-class ViewReference(Identifier):
+class TypedViewReference(Identifier):
     type: Literal["view"] = "view"
     space: str
     external_id: str
@@ -88,108 +97,139 @@ class ViewReference(Identifier):
 #######################################################
 
 
-class InstanceRequestResource(RequestResource):
-    """This is a base class for resources that are Instances."""
+class WrappedInstanceRequest(RequestResource, ABC):
+    """This is a base class for resources that are Instances.
+    It is used to define resources that are
+    """
 
-    VIEW_ID: ClassVar[ViewReference]
+    VIEW_ID: ClassVar[TypedViewReference]
     instance_type: InstanceType
     space: str
     external_id: str
+    existing_version: int | None = None
 
-    def as_id(self) -> TypedInstanceIdentifier:
-        return TypedInstanceIdentifier(
-            instance_type=self.instance_type,
-            space=self.space,
-            external_id=self.external_id,
-        )
-
-    def as_request_item(self) -> "InstanceRequestItem":
-        return InstanceRequestItem(
-            instance_type=self.instance_type,
-            space=self.space,
-            external_id=self.external_id,
-            sources=[InstanceSource(source=self.VIEW_ID, resource=self)],
-        )
-
-
-class InstanceSource(BaseModelObject):
-    source: ViewReference
-    resource: InstanceRequestResource
+    def dump(
+        self, camel_case: bool = True, exclude_extra: bool = False, context: Literal["api", "toolkit"] = "api"
+    ) -> dict[str, Any]:
+        """Dump the resource to a dictionary.
 
-    @model_serializer(mode="plain")
-    def serialize_resource(self) -> dict[str, Any]:
-        properties: dict[str, JsonValue] = {}
-        for field_id, field in type(self.resource).model_fields.items():
-            if field_id in InstanceRequestResource.model_fields:
-                # Skip space, external_id, instance_type
-                continue
-            key = field.alias or field_id
-            properties[key] = self._serialize_property(getattr(self.resource, field_id))
+        Args:
+            camel_case (bool): Whether to use camelCase for the keys. Default is True.
+            exclude_extra (bool): Whether to exclude extra fields not defined in the model. Default is False.
+            context (Literal["api", "toolkit"]): The context in which the dump is used. Default is "api".
 
+        """
+        exclude: set[str] = set()
+        if exclude_extra:
+            exclude |= set(self.__pydantic_extra__) if self.__pydantic_extra__ else set()
+        if context == "api":
+            exclude.update({"existing_version", "instance_type", "space", "external_id"})
+        dumped = self.model_dump(mode="json", by_alias=camel_case, exclude_unset=True, exclude=exclude)
+        if context == "toolkit":
+            return dumped
         return {
-            "source": self.source.model_dump(by_alias=True),
-            "properties": properties,
+            "instanceType": self.instance_type,
+            "space": self.space,
+            "externalId": self.external_id,
+            "sources": [
+                {
+                    "source": self.VIEW_ID.dump(camel_case=camel_case, include_type=True),
+                    "properties": dumped,
+                }
+            ],
         }
 
-    @classmethod
-    def _serialize_property(cls, value: Any) -> JsonValue:
-        """Handles serialization of direct relations."""
-        if isinstance(value, InstanceRequestResource):
-            return {"space": value.space, "externalId": value.external_id}
-        elif isinstance(value, list):
-            return [cls._serialize_property(v) for v in value]
-        return value
 
+T_WrappedInstanceRequest = TypeVar("T_WrappedInstanceRequest", bound=WrappedInstanceRequest)
 
-class InstanceRequestItem(RequestResource):
-    model_config = ConfigDict(populate_by_name=True)
-    instance_type: InstanceType
-    space: str
-    external_id: str
-    existing_version: int | None = None
-    sources: list[InstanceSource] | None = None
-
-    def as_id(self) -> TypedInstanceIdentifier:
-        return TypedInstanceIdentifier(
-            instance_type=self.instance_type,
-            space=self.space,
-            external_id=self.external_id,
-        )
 
-
-class InstanceResponseItem(BaseModelObject):
+class WrappedInstanceResponse(ResponseResource[T_WrappedInstanceRequest], ABC):
+    VIEW_ID: ClassVar[TypedViewReference]
     instance_type: InstanceType
     space: str
     external_id: str
+
     version: int
-    type: TypedInstanceIdentifier | None = None
     created_time: int
     last_updated_time: int
     deleted_time: int | None = None
-    properties: dict[str, dict[str, dict[str, JsonValue]]] | None = None
 
-    def get_properties_for_source(
-        self, source: ViewReference, include_identifier: bool = False
-    ) -> dict[str, JsonValue]:
-        output: dict[str, JsonValue] = (
-            {"space": self.space, "externalId": self.external_id} if include_identifier else {}
-        )
-        if not self.properties:
-            return output
-        if source.space not in self.properties:
-            return output
-        space_properties = self.properties[source.space]
-        view_version = f"{source.external_id}/{source.version}"
-        output.update(space_properties.get(view_version, {}))
-        return output
+    @model_validator(mode="before")
+    def move_properties(cls, values: dict[str, Any]) -> dict[str, Any]:
+        """Move properties from sources to the top level."""
+        return move_properties(values, cls.VIEW_ID)
 
-    def as_id(self) -> TypedInstanceIdentifier:
-        return TypedInstanceIdentifier(
+
+def move_properties(values: dict[str, Any], view_id: TypedViewReference) -> dict[str, Any]:
+    """Help function to move properties from properties.space.externalId/version to the top level.
+
+    It is used in WrappedInstanceResponse to move properties from the response to the top level.
+    """
+    if "properties" not in values:
+        return values
+    values_copy = dict(values)
+    properties = values_copy.pop("properties")
+    if not isinstance(properties, dict) or view_id.space not in properties:
+        return values
+    view_properties = properties.pop(view_id.space)
+    identifier = f"{view_id.external_id}/{view_id.version}"
+    if not isinstance(view_properties, dict) or identifier not in view_properties:
+        return values
+    source_properties = view_properties.pop(identifier)
+    values_copy.update(source_properties)
+    return values_copy
+
+
+T_WrappedInstanceResponse = TypeVar("T_WrappedInstanceResponse", bound=WrappedInstanceResponse)
+
+
+class WrappedInstanceListRequest(RequestResource, ABC):
+    VIEW_ID: ClassVar[TypedViewReference]
+    instance_type: Literal["node"] = "node"
+    space: str
+    external_id: str
+
+    @abstractmethod
+    def dump_instances(self) -> list[dict[str, Any]]:
+        """Dumps the object to a list of instance request dictionaries."""
+        raise NotImplementedError()
+
+    def as_id(self) -> TypedNodeIdentifier:
+        return TypedNodeIdentifier(
             instance_type=self.instance_type,
             space=self.space,
             external_id=self.external_id,
         )
 
+    @abstractmethod
+    def as_ids(self) -> list[TypedInstanceIdentifier]:
+        """Convert the response to a list of typed instance identifiers."""
+        raise NotImplementedError()
+
+
+T_InstancesListRequest = TypeVar("T_InstancesListRequest", bound=WrappedInstanceListRequest)
+
+
+class WrappedInstanceListResponse(ResponseResource[T_InstancesListRequest], ABC):
+    VIEW_ID: ClassVar[TypedViewReference]
+    instance_type: Literal["node"] = "node"
+    space: str
+    external_id: str
+
+    @model_validator(mode="before")
+    @classmethod
+    def move_properties(cls, values: dict[str, Any]) -> dict[str, Any]:
+        """Move properties from sources to the top level."""
+        return move_properties(values, cls.VIEW_ID)
+
+    @abstractmethod
+    def as_ids(self) -> list[TypedInstanceIdentifier]:
+        """Convert the response to a list of typed instance identifiers."""
+        raise NotImplementedError()
+
+
+T_InstancesListResponse = TypeVar("T_InstancesListResponse", bound=WrappedInstanceListResponse)
+
 
 class NodeReference(BaseModelObject):
     space: str