cognite-neat 0.88.2__py3-none-any.whl → 0.89.0__py3-none-any.whl

cognite/neat/rules/models/dms/_rules_input.py

@@ -1,24 +1,29 @@
-from collections.abc import Sequence
+import re
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Literal, cast, overload
+from typing import Any, Literal
 
-from cognite.neat.rules.models._base import DataModelType, ExtensionCategory, SchemaCompleteness, _add_alias
+from cognite.client import data_modeling as dm
+
+from cognite.neat.rules.models._base_input import InputComponent, InputRules
 from cognite.neat.rules.models.data_types import DataType
 from cognite.neat.rules.models.entities import (
     ClassEntity,
     ContainerEntity,
+    DMSNodeEntity,
     DMSUnknownEntity,
-    Unknown,
+    EdgeEntity,
+    ReverseConnectionEntity,
     ViewEntity,
-    ViewPropertyEntity,
+    load_connection,
+    load_dms_value_type,
 )
 
-from ._rules import DMSContainer, DMSMetadata, DMSProperty, DMSRules, DMSView
+from ._rules import _DEFAULT_VERSION, DMSContainer, DMSEnum, DMSMetadata, DMSNode, DMSProperty, DMSRules, DMSView
 
 
 @dataclass
-class DMSMetadataInput:
+class DMSInputMetadata(InputComponent[DMSMetadata]):
     schema_: Literal["complete", "partial", "extended"]
     space: str
     external_id: str
@@ -32,51 +37,57 @@ class DMSMetadataInput:
     updated: datetime | str | None = None
 
     @classmethod
-    def load(cls, data: dict[str, Any] | None) -> "DMSMetadataInput | None":
-        if data is None:
-            return None
-        _add_alias(data, DMSMetadata)
+    def _get_verified_cls(cls) -> type[DMSMetadata]:
+        return DMSMetadata
+
+    def dump(self) -> dict[str, Any]:  # type: ignore[override]
+        output = super().dump()
+        if self.created is None:
+            output["created"] = datetime.now()
+        if self.updated is None:
+            output["updated"] = datetime.now()
+        return output
+
+    @classmethod
+    def from_data_model(cls, data_model: dm.DataModelApply, has_reference: bool) -> "DMSInputMetadata":
+        description, creator = cls._get_description_and_creator(data_model.description)
         return cls(
-            schema_=data.get("schema_"),  # type: ignore[arg-type]
-            space=data.get("space"),  # type: ignore[arg-type]
-            external_id=data.get("external_id"),  # type: ignore[arg-type]
-            creator=data.get("creator"),  # type: ignore[arg-type]
-            version=data.get("version"),  # type: ignore[arg-type]
-            # safeguard from empty cell, i.e. if key provided by value None
-            extension=data.get("extension", "addition") or "addition",
-            data_model_type=data.get("data_model_type", "solution") or "solution",
-            name=data.get("name"),
-            description=data.get("description"),
-            created=data.get("created"),
-            updated=data.get("updated"),
+            schema_="complete",
+            data_model_type="solution" if has_reference else "enterprise",
+            space=data_model.space,
+            name=data_model.name or None,
+            description=description,
+            external_id=data_model.external_id,
+            version=data_model.version,
+            creator=",".join(creator),
+            created=datetime.now(),
+            updated=datetime.now(),
         )
 
-    def dump(self) -> dict[str, Any]:
-        return dict(
-            schema=SchemaCompleteness(self.schema_),
-            extension=ExtensionCategory(self.extension),
-            space=self.space,
-            externalId=self.external_id,
-            dataModelType=DataModelType(self.data_model_type),
-            creator=self.creator,
-            version=self.version,
-            name=self.name,
-            description=self.description,
-            created=self.created or datetime.now(),
-            updated=self.updated or datetime.now(),
-        )
+    @classmethod
+    def _get_description_and_creator(cls, description_raw: str | None) -> tuple[str | None, list[str]]:
+        if description_raw and (description_match := re.search(r"Creator: (.+)", description_raw)):
+            creator = description_match.group(1).split(", ")
+            description = description_raw.replace(description_match.string, "").strip() or None
+        elif description_raw:
+            creator = ["MISSING"]
+            description = description_raw
+        else:
+            creator = ["MISSING"]
+            description = None
+        return description, creator
 
 
 @dataclass
-class DMSPropertyInput:
+class DMSInputProperty(InputComponent[DMSProperty]):
     view: str
     view_property: str | None
-    value_type: str
-    property_: str | None
+    value_type: str | DataType | ViewEntity | DMSUnknownEntity
+    property_: str | None = None
     class_: str | None = None
     name: str | None = None
     description: str | None = None
-    connection: Literal["direct", "edge", "reverse"] | None = None
+    connection: Literal["direct"] | ReverseConnectionEntity | EdgeEntity | str | None = None
     nullable: bool | None = None
     immutable: bool | None = None
     is_list: bool | None = None
@@ -88,153 +99,76 @@ class DMSPropertyInput:
     constraint: str | list[str] | None = None
 
     @classmethod
-    @overload
-    def load(cls, data: None) -> None: ...
-
-    @classmethod
-    @overload
-    def load(cls, data: dict[str, Any]) -> "DMSPropertyInput": ...
-
-    @classmethod
-    @overload
-    def load(cls, data: list[dict[str, Any]]) -> list["DMSPropertyInput"]: ...
-
-    @classmethod
-    def load(
-        cls, data: dict[str, Any] | list[dict[str, Any]] | None
-    ) -> "DMSPropertyInput | list[DMSPropertyInput] | None":
-        if data is None:
-            return None
-        if isinstance(data, list) or (isinstance(data, dict) and isinstance(data.get("data"), list)):
-            items = cast(list[dict[str, Any]], data.get("data") if isinstance(data, dict) else data)
-            return [loaded for item in items if (loaded := cls.load(item)) is not None]
-
-        _add_alias(data, DMSProperty)
-        return cls(
-            view=data.get("view"),  # type: ignore[arg-type]
-            view_property=data.get("view_property"),  # type: ignore[arg-type]
-            value_type=data.get("value_type"),  # type: ignore[arg-type]
-            property_=data.get("property_"),
-            class_=data.get("class_"),
-            name=data.get("name"),
-            description=data.get("description"),
-            connection=data.get("connection"),
-            nullable=data.get("nullable"),
-            immutable=data.get("immutable"),
-            is_list=data.get("is_list"),
-            default=data.get("default"),
-            reference=data.get("reference"),
-            container=data.get("container"),
-            container_property=data.get("container_property"),
-            index=data.get("index"),
-            constraint=data.get("constraint"),
+    def _get_verified_cls(cls) -> type[DMSProperty]:
+        return DMSProperty
+
+    def dump(self, default_space: str, default_version: str) -> dict[str, Any]:  # type: ignore[override]
+        output = super().dump()
+        output["View"] = ViewEntity.load(self.view, space=default_space, version=default_version)
+        output["Value Type"] = load_dms_value_type(self.value_type, default_space, default_version)
+        output["Connection"] = load_connection(self.connection, default_space, default_version)
+        output["Property (linage)"] = self.property_ or self.view_property
+        output["Class (linage)"] = (
+            ClassEntity.load(self.class_ or self.view, prefix=default_space, version=default_version)
+            if self.class_ or self.view
+            else None
         )
-
-    def dump(self, default_space: str, default_version: str) -> dict[str, Any]:
-        value_type: DataType | ViewPropertyEntity | ViewEntity | DMSUnknownEntity
-        if DataType.is_data_type(self.value_type):
-            value_type = DataType.load(self.value_type)
-        elif self.value_type == str(Unknown):
-            value_type = DMSUnknownEntity()
-        else:
-            try:
-                value_type = ViewPropertyEntity.load(self.value_type, space=default_space, version=default_version)
-            except ValueError:
-                value_type = ViewEntity.load(self.value_type, space=default_space, version=default_version)
-
-        return {
-            "View": ViewEntity.load(self.view, space=default_space, version=default_version),
-            "View Property": self.view_property,
-            "Value Type": value_type,
-            "Property (linage)": self.property_ or self.view_property,
-            "Class (linage)": (
-                ClassEntity.load(self.class_ or self.view, prefix=default_space, version=default_version)
-                if self.class_ or self.view
-                else None
-            ),
-            "Name": self.name,
-            "Description": self.description,
-            "Connection": self.connection,
-            "Nullable": self.nullable,
-            "Immutable": self.immutable,
-            "Is List": self.is_list,
-            "Default": self.default,
-            "Reference": self.reference,
-            "Container": (
-                ContainerEntity.load(self.container, space=default_space, version=default_version)
-                if self.container
-                else None
-            ),
-            "Container Property": self.container_property,
-            "Index": self.index,
-            "Constraint": self.constraint,
-        }
+        output["Container"] = (
+            ContainerEntity.load(self.container, space=default_space, version=default_version)
+            if self.container
+            else None
+        )
+        return output
 
 
 @dataclass
-class DMSContainerInput:
+class DMSInputContainer(InputComponent[DMSContainer]):
     container: str
     class_: str | None = None
     name: str | None = None
     description: str | None = None
     reference: str | None = None
     constraint: str | None = None
+    used_for: Literal["node", "edge", "all"] | None = None
 
     @classmethod
-    @overload
-    def load(cls, data: None) -> None: ...
-
-    @classmethod
-    @overload
-    def load(cls, data: dict[str, Any]) -> "DMSContainerInput": ...
+    def _get_verified_cls(cls) -> type[DMSContainer]:
+        return DMSContainer
 
-    @classmethod
-    @overload
-    def load(cls, data: list[dict[str, Any]]) -> list["DMSContainerInput"]: ...
+    def dump(self, default_space: str) -> dict[str, Any]:  # type: ignore[override]
+        output = super().dump()
+        container = ContainerEntity.load(self.container, space=default_space)
+        output["Container"] = container
+        output["Class (linage)"] = (
+            ClassEntity.load(self.class_, prefix=default_space) if self.class_ else container.as_class()
+        )
+        output["Constraint"] = (
+            [ContainerEntity.load(constraint.strip(), space=default_space) for constraint in self.constraint.split(",")]
+            if self.constraint
+            else None
+        )
+        return output
 
     @classmethod
-    def load(
-        cls, data: dict[str, Any] | list[dict[str, Any]] | None
-    ) -> "DMSContainerInput | list[DMSContainerInput] | None":
-        if data is None:
-            return None
-        if isinstance(data, list) or (isinstance(data, dict) and isinstance(data.get("data"), list)):
-            items = cast(list[dict[str, Any]], data.get("data") if isinstance(data, dict) else data)
-            return [loaded for item in items if (loaded := cls.load(item)) is not None]
-
-        _add_alias(data, DMSContainer)
+    def from_container(cls, container: dm.ContainerApply) -> "DMSInputContainer":
+        constraints: list[str] = []
+        for _, constraint_obj in (container.constraints or {}).items():
+            if isinstance(constraint_obj, dm.RequiresConstraint):
+                constraints.append(str(ContainerEntity.from_id(constraint_obj.require)))
+            # UniquenessConstraint it handled in the properties
+        container_entity = ContainerEntity.from_id(container.as_id())
         return cls(
-            container=data.get("container"),  # type: ignore[arg-type]
-            class_=data.get("class_"),
-            name=data.get("name"),
-            description=data.get("description"),
-            reference=data.get("reference"),
-            constraint=data.get("constraint"),
+            class_=str(container_entity.as_class()),
+            container=str(container_entity),
+            name=container.name or None,
+            description=container.description,
+            constraint=", ".join(constraints) or None,
+            used_for=container.used_for,
         )
 
-    def dump(self, default_space: str) -> dict[str, Any]:
-        container = ContainerEntity.load(self.container, space=default_space)
-        return {
-            "Container": container,
-            "Class (linage)": (
-                ClassEntity.load(self.class_, prefix=default_space) if self.class_ else container.as_class()
-            ),
-            "Name": self.name,
-            "Description": self.description,
-            "Reference": self.reference,
-            "Constraint": (
-                [
-                    ContainerEntity.load(constraint.strip(), space=default_space)
-                    for constraint in self.constraint.split(",")
-                ]
-                if self.constraint
-                else None
-            ),
-        }
-
 
 @dataclass
-class DMSViewInput:
+class DMSInputView(InputComponent[DMSView]):
     view: str
     class_: str | None = None
     name: str | None = None
@@ -245,117 +179,115 @@ class DMSViewInput:
     in_model: bool = True
 
     @classmethod
-    @overload
-    def load(cls, data: None) -> None: ...
-
-    @classmethod
-    @overload
-    def load(cls, data: dict[str, Any]) -> "DMSViewInput": ...
+    def _get_verified_cls(cls) -> type[DMSView]:
+        return DMSView
 
-    @classmethod
-    @overload
-    def load(cls, data: list[dict[str, Any]]) -> list["DMSViewInput"]: ...
+    def dump(self, default_space: str, default_version: str) -> dict[str, Any]:  # type: ignore[override]
+        output = super().dump()
+        view = ViewEntity.load(self.view, space=default_space, version=default_version)
+        output["View"] = view
+        output["Class (linage)"] = (
+            ClassEntity.load(self.class_, prefix=default_space, version=default_version)
+            if self.class_
+            else view.as_class()
+        )
+        output["Implements"] = (
+            [
+                ViewEntity.load(implement, space=default_space, version=default_version)
+                for implement in self.implements.split(",")
+            ]
+            if self.implements
+            else None
+        )
+        return output
 
     @classmethod
-    def load(cls, data: dict[str, Any] | list[dict[str, Any]] | None) -> "DMSViewInput | list[DMSViewInput] | None":
-        if data is None:
-            return None
-        if isinstance(data, list) or (isinstance(data, dict) and isinstance(data.get("data"), list)):
-            items = cast(list[dict[str, Any]], data.get("data") if isinstance(data, dict) else data)
-            return [loaded for item in items if (loaded := cls.load(item)) is not None]
-        _add_alias(data, DMSView)
+    def from_view(cls, view: dm.ViewApply, in_model: bool) -> "DMSInputView":
+        view_entity = ViewEntity.from_id(view.as_id())
+        class_entity = view_entity.as_class(skip_version=True)
 
         return cls(
-            view=data.get("view"),  # type: ignore[arg-type]
-            class_=data.get("class"),
-            name=data.get("name"),
-            description=data.get("description"),
-            implements=data.get("implements"),
-            reference=data.get("reference"),
-            filter_=data.get("filter_"),
-            in_model=data.get("in_model", True),
+            class_=str(class_entity),
+            view=str(view_entity),
+            description=view.description,
+            name=view.name,
+            implements=", ".join([str(ViewEntity.from_id(parent, _DEFAULT_VERSION)) for parent in view.implements])
+            or None,
+            in_model=in_model,
         )
 
-    def dump(self, default_space: str, default_version: str) -> dict[str, Any]:
-        view = ViewEntity.load(self.view, space=default_space, version=default_version)
-        return {
-            "View": view,
-            "Class (linage)": (
-                ClassEntity.load(self.class_, prefix=default_space, version=default_version)
-                if self.class_
-                else view.as_class()
-            ),
-            "Name": self.name,
-            "Description": self.description,
-            "Implements": (
-                [
-                    ViewEntity.load(implement, space=default_space, version=default_version)
-                    for implement in self.implements.split(",")
-                ]
-                if self.implements
-                else None
-            ),
-            "Reference": self.reference,
-            "Filter": self.filter_,
-            "In Model": self.in_model,
-        }
-
 
 @dataclass
-class DMSRulesInput:
-    metadata: DMSMetadataInput
-    properties: Sequence[DMSPropertyInput]
-    views: Sequence[DMSViewInput]
-    containers: Sequence[DMSContainerInput] | None = None
-    last: "DMSRulesInput | DMSRules | None" = None
-    reference: "DMSRulesInput | DMSRules | None" = None
+class DMSInputNode(InputComponent[DMSNode]):
+    node: str
+    usage: Literal["type", "collocation"]
+    name: str | None = None
+    description: str | None = None
 
     @classmethod
-    @overload
-    def load(cls, data: dict[str, Any]) -> "DMSRulesInput": ...
+    def _get_verified_cls(cls) -> type[DMSNode]:
+        return DMSNode
 
     @classmethod
-    @overload
-    def load(cls, data: None) -> None: ...
+    def from_node_type(cls, node_type: dm.NodeApply) -> "DMSInputNode":
+        return cls(node=f"{node_type.space}:{node_type.external_id}", usage="type")
+
+    def dump(self, default_space: str, **_) -> dict[str, Any]:  # type: ignore[override]
+        output = super().dump()
+        output["Node"] = DMSNodeEntity.load(self.node, space=default_space)
+        return output
+
+
+@dataclass
+class DMSInputEnum(InputComponent[DMSEnum]):
+    collection: str
+    value: str
+    name: str | None = None
+    description: str | None = None
 
     @classmethod
-    def load(cls, data: dict | None) -> "DMSRulesInput | None":
-        if data is None:
-            return None
-        _add_alias(data, DMSRules)
-        return cls(
-            metadata=DMSMetadataInput.load(data.get("metadata")),  # type: ignore[arg-type]
-            properties=DMSPropertyInput.load(data.get("properties")),  # type: ignore[arg-type]
-            views=DMSViewInput.load(data.get("views")),  # type: ignore[arg-type]
-            containers=DMSContainerInput.load(data.get("containers")) or [],
-            last=DMSRulesInput.load(data.get("last")),
-            reference=DMSRulesInput.load(data.get("reference")),
-        )
+    def _get_verified_cls(cls) -> type[DMSEnum]:
+        return DMSEnum
+
 
-    def as_rules(self) -> DMSRules:
-        return DMSRules.model_validate(self.dump())
+@dataclass
+class DMSInputRules(InputRules[DMSRules]):
+    metadata: DMSInputMetadata
+    properties: list[DMSInputProperty]
+    views: list[DMSInputView]
+    containers: list[DMSInputContainer] | None = None
+    enum: list[DMSInputEnum] | None = None
+    nodes: list[DMSInputNode] | None = None
+    last: "DMSInputRules | None" = None
+    reference: "DMSInputRules | None" = None
+
+    @classmethod
+    def _get_verified_cls(cls) -> type[DMSRules]:
+        return DMSRules
 
     def dump(self) -> dict[str, Any]:
         default_space = self.metadata.space
         default_version = self.metadata.version
         reference: dict[str, Any] | None = None
-        if isinstance(self.reference, DMSRulesInput):
+        if isinstance(self.reference, DMSInputRules):
             reference = self.reference.dump()
         elif isinstance(self.reference, DMSRules):
             # We need to load through the DMSRulesInput to set the correct default space and version
-            reference = DMSRulesInput.load(self.reference.model_dump()).dump()
+            reference = DMSInputRules.load(self.reference.model_dump()).dump()
         last: dict[str, Any] | None = None
-        if isinstance(self.last, DMSRulesInput):
+        if isinstance(self.last, DMSInputRules):
             last = self.last.dump()
         elif isinstance(self.last, DMSRules):
             # We need to load through the DMSRulesInput to set the correct default space and version
-            last = DMSRulesInput.load(self.last.model_dump()).dump()
-
-        return dict(
-            Metadata=self.metadata.dump(),
-            Properties=[prop.dump(default_space, default_version) for prop in self.properties],
-            Views=[view.dump(default_space, default_version) for view in self.views],
-            Containers=[container.dump(default_space) for container in self.containers or []] or None,
-            Last=last,
-            Reference=reference,
-        )
+            last = DMSInputRules.load(self.last.model_dump()).dump()
+
+        return {
+            "Metadata": self.metadata.dump(),
+            "Properties": [prop.dump(default_space, default_version) for prop in self.properties],
+            "Views": [view.dump(default_space, default_version) for view in self.views],
+            "Containers": [container.dump(default_space) for container in self.containers or []] or None,
+            "Enum": [enum.dump() for enum in self.enum or []] or None,
+            "Nodes": [node_type.dump(default_space) for node_type in self.nodes or []] or None,
+            "Last": last,
+            "Reference": reference,
+        }