cognite-neat 0.88.3__py3-none-any.whl → 0.90.0__py3-none-any.whl

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

@@ -1,10 +1,11 @@
 import warnings
 from collections import defaultdict
-from collections.abc import Sequence
+from collections.abc import Collection, Sequence
 from typing import Any, cast
 
 from cognite.client.data_classes import data_modeling as dm
 from cognite.client.data_classes.data_modeling.containers import BTreeIndex
+from cognite.client.data_classes.data_modeling.data_types import EnumValue as DMSEnumValue
 from cognite.client.data_classes.data_modeling.data_types import ListablePropertyType
 from cognite.client.data_classes.data_modeling.views import (
     SingleEdgeConnectionApply,
@@ -12,6 +13,7 @@ from cognite.client.data_classes.data_modeling.views import (
     ViewPropertyApply,
 )
 
+from cognite.neat.issues.errors import NeatTypeError, ResourceNotFoundError
 from cognite.neat.issues.warnings import NotSupportedWarning, PropertyNotFoundWarning
 from cognite.neat.issues.warnings.user_modeling import (
     EmptyContainerWarning,
@@ -19,20 +21,25 @@ from cognite.neat.issues.warnings.user_modeling import (
     HasDataFilterOnViewWithReferencesWarning,
     NodeTypeFilterOnParentViewWarning,
 )
-from cognite.neat.rules.models._base import DataModelType, ExtensionCategory, SchemaCompleteness
-from cognite.neat.rules.models.data_types import DataType
+from cognite.neat.rules.models._base_rules import DataModelType, ExtensionCategory, SchemaCompleteness
+from cognite.neat.rules.models.data_types import DataType, Double, Enum, Float
 from cognite.neat.rules.models.entities import (
+    ClassEntity,
     ContainerEntity,
+    DMSFilter,
     DMSNodeEntity,
     DMSUnknownEntity,
+    EdgeEntity,
+    HasDataFilter,
+    NodeTypeFilter,
     ReferenceEntity,
+    ReverseConnectionEntity,
+    UnitEntity,
     ViewEntity,
-    ViewPropertyEntity,
 )
-from cognite.neat.rules.models.wrapped_entities import DMSFilter, HasDataFilter, NodeTypeFilter
 from cognite.neat.utils.cdf.data_classes import ContainerApplyDict, NodeApplyDict, SpaceApplyDict, ViewApplyDict
 
-from ._rules import DMSMetadata, DMSProperty, DMSRules, DMSView
+from ._rules import DMSEnum, DMSMetadata, DMSProperty, DMSRules, DMSView
 from ._schema import DMSSchema, PipelineSchema
 
 
@@ -111,9 +118,16 @@ class _DMSExporter:
             ]
             self._update_with_properties(selected_properties, container_properties_by_id, None)
 
-        containers = self._create_containers(container_properties_by_id)
+        containers = self._create_containers(container_properties_by_id, rules.enum)  # type: ignore[arg-type]
 
-        views, node_types = self._create_views_with_node_types(view_properties_by_id)
+        views, view_node_type_filters = self._create_views_with_node_types(view_properties_by_id)
+        if rules.nodes:
+            node_types = NodeApplyDict(
+                [node.as_node() for node in rules.nodes]
+                + [dm.NodeApply(node.space, node.external_id) for node in view_node_type_filters]
+            )
+        else:
+            node_types = NodeApplyDict([dm.NodeApply(node.space, node.external_id) for node in view_node_type_filters])
 
         last_schema: DMSSchema | None = None
         if self.rules.last:
@@ -184,7 +198,7 @@ class _DMSExporter:
     def _create_views_with_node_types(
         self,
         view_properties_by_id: dict[dm.ViewId, list[DMSProperty]],
-    ) -> tuple[ViewApplyDict, NodeApplyDict]:
+    ) -> tuple[ViewApplyDict, set[dm.NodeId]]:
         input_views = list(self.rules.views)
         if self.rules.last:
             existing = {view.view.as_id() for view in input_views}
@@ -245,17 +259,19 @@ class _DMSExporter:
                 # as they are expected for the solution model.
                 unique_node_types.add(dm.NodeId(space=view.space, external_id=view.external_id))
 
-        return views, NodeApplyDict(
-            [dm.NodeApply(space=node.space, external_id=node.external_id) for node in unique_node_types]
-        )
+        return views, unique_node_types
 
     @classmethod
     def _create_edge_type_from_prop(cls, prop: DMSProperty) -> dm.DirectRelationReference:
-        if isinstance(prop.reference, ReferenceEntity):
+        if isinstance(prop.connection, EdgeEntity) and prop.connection.edge_type is not None:
+            return prop.connection.edge_type.as_reference()
+        elif isinstance(prop.reference, ReferenceEntity):
             ref_view_prop = prop.reference.as_view_property_id()
             return cls._create_edge_type_from_view_id(cast(dm.ViewId, ref_view_prop.source), ref_view_prop.property)
-        else:
+        elif isinstance(prop.value_type, ViewEntity):
             return cls._create_edge_type_from_view_id(prop.view.as_id(), prop.view_property)
+        else:
+            raise NeatTypeError(f"Invalid valueType {prop.value_type!r}")
 
     @staticmethod
     def _create_edge_type_from_view_id(view_id: dm.ViewId, property_: str) -> dm.DirectRelationReference:
@@ -268,7 +284,12 @@ class _DMSExporter:
     def _create_containers(
         self,
         container_properties_by_id: dict[dm.ContainerId, list[DMSProperty]],
+        enum: Collection[DMSEnum] | None,
     ) -> ContainerApplyDict:
+        enum_values_by_collection: dict[ClassEntity, list[DMSEnum]] = defaultdict(list)
+        for enum_value in enum or []:
+            enum_values_by_collection[enum_value.collection].append(enum_value)
+
         containers = list(self.rules.containers or [])
         if self.rules.last:
             existing = {container.container.as_id() for container in containers}
@@ -297,17 +318,34 @@ class _DMSExporter:
                     type_cls = prop.value_type.dms
                 else:
                     type_cls = dm.DirectRelation
-                type_: dm.PropertyType
+
+                args: dict[str, Any] = {}
                 if issubclass(type_cls, ListablePropertyType):
-                    type_ = type_cls(is_list=prop.is_list or False)
-                else:
-                    type_ = type_cls()
+                    args["is_list"] = prop.is_list or False
+                if isinstance(prop.value_type, Double | Float) and isinstance(prop.value_type.unit, UnitEntity):
+                    args["unit"] = prop.value_type.unit.as_reference()
+                if isinstance(prop.value_type, Enum):
+                    if prop.value_type.collection not in enum_values_by_collection:
+                        raise ResourceNotFoundError(
+                            prop.value_type.collection, "enum collection", prop.container, "container"
+                        )
+                    args["unknown_value"] = prop.value_type.unknown_value
+                    args["values"] = {
+                        value.value: DMSEnumValue(
+                            name=value.name,
+                            description=value.description,
+                        )
+                        for value in enum_values_by_collection[prop.value_type.collection]
+                    }
+
+                type_ = type_cls(**args)
                 container.properties[prop.container_property] = dm.ContainerProperty(
                     type=type_,
                     # If not set, nullable is True and immutable is False
                     nullable=prop.nullable if prop.nullable is not None else True,
                     immutable=prop.immutable if prop.immutable is not None else False,
-                    default_value=prop.default,
+                    # Guarding against default value being set for connection properties
+                    default_value=prop.default if not prop.connection else None,
                     name=prop.name,
                     description=prop.description,
                 )
@@ -463,7 +501,7 @@ class _DMSExporter:
                 description=prop.description,
                 **extra_args,
             )
-        elif prop.connection == "edge":
+        elif isinstance(prop.connection, EdgeEntity):
             if isinstance(prop.value_type, ViewEntity):
                 source_view_id = prop.value_type.as_id()
             else:
@@ -472,6 +510,9 @@ class _DMSExporter:
                     "If this error occurs it is a bug in NEAT, please report"
                     f"Debug Info, Invalid valueType edge: {prop.model_dump_json()}"
                 )
+            edge_source: dm.ViewId | None = None
+            if prop.connection.properties is not None:
+                edge_source = prop.connection.properties.as_id()
             edge_cls: type[dm.EdgeConnectionApply] = dm.MultiEdgeConnectionApply
             # If is_list is not set, we default to a MultiEdgeConnection
             if prop.is_list is False:
@@ -483,13 +524,11 @@ class _DMSExporter:
                 direction="outwards",
                 name=prop.name,
                 description=prop.description,
+                edge_source=edge_source,
             )
-        elif prop.connection == "reverse":
-            reverse_prop_id: str | None = None
-            if isinstance(prop.value_type, ViewPropertyEntity):
-                source_view_id = prop.value_type.as_view_id()
-                reverse_prop_id = prop.value_type.property_
-            elif isinstance(prop.value_type, ViewEntity):
+        elif isinstance(prop.connection, ReverseConnectionEntity):
+            reverse_prop_id = prop.connection.property_
+            if isinstance(prop.value_type, ViewEntity):
                 source_view_id = prop.value_type.as_id()
             else:
                 # Should have been validated.
@@ -497,16 +536,17 @@ class _DMSExporter:
                     "If this error occurs it is a bug in NEAT, please report"
                     f"Debug Info, Invalid valueType reverse connection: {prop.model_dump_json()}"
                 )
-            reverse_prop: DMSProperty | None = None
-            if reverse_prop_id is not None:
-                reverse_prop = next(
-                    (
-                        prop
-                        for prop in view_properties_by_id.get(source_view_id, [])
-                        if prop.property_ == reverse_prop_id
-                    ),
-                    None,
-                )
+            edge_source = None
+            reverse_prop = next(
+                (prop for prop in view_properties_by_id.get(source_view_id, []) if prop.property_ == reverse_prop_id),
+                None,
+            )
+            if (
+                reverse_prop
+                and isinstance(reverse_prop.connection, EdgeEntity)
+                and reverse_prop.connection.properties is not None
+            ):
+                edge_source = reverse_prop.connection.properties.as_id()
 
             if reverse_prop is None:
                 warnings.warn(
@@ -520,7 +560,7 @@ class _DMSExporter:
                     stacklevel=2,
                 )
 
-            if reverse_prop is None or reverse_prop.connection == "edge":
+            if reverse_prop is None or isinstance(reverse_prop.connection, EdgeEntity):
                 inwards_edge_cls = (
                     dm.MultiEdgeConnectionApply if prop.is_list in [True, None] else SingleEdgeConnectionApply
                 )
@@ -530,10 +570,13 @@ class _DMSExporter:
                     name=prop.name,
                     description=prop.description,
                     direction="inwards",
+                    edge_source=edge_source,
                 )
-            elif reverse_prop_id and reverse_prop and reverse_prop.connection == "direct":
+            elif reverse_prop and reverse_prop.connection == "direct":
                 reverse_direct_cls = (
-                    dm.MultiReverseDirectRelationApply if prop.is_list is True else SingleReverseDirectRelationApply
+                    dm.MultiReverseDirectRelationApply
+                    if prop.is_list in [True, None]
+                    else SingleReverseDirectRelationApply
                 )
                 return reverse_direct_cls(
                     source=source_view_id,

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

@@ -1,8 +1,6 @@
 import math
-import re
 import sys
 import warnings
-from collections import defaultdict
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, ClassVar, Literal
 
@@ -16,7 +14,7 @@ from cognite.neat.issues.warnings import (
     PrincipleMatchingSpaceAndVersionWarning,
     PrincipleSolutionBuildsOnEnterpriseWarning,
 )
-from cognite.neat.rules.models._base import (
+from cognite.neat.rules.models._base_rules import (
     BaseMetadata,
     BaseRules,
     DataModelType,
@@ -33,24 +31,27 @@ from cognite.neat.rules.models._types import (
     VersionType,
 )
 from cognite.neat.rules.models.data_types import DataType
-from cognite.neat.rules.models.domain import DomainRules
 from cognite.neat.rules.models.entities import (
     ClassEntity,
     ContainerEntity,
     ContainerEntityList,
+    DMSNodeEntity,
     DMSUnknownEntity,
+    EdgeEntity,
+    HasDataFilter,
+    NodeTypeFilter,
+    RawFilter,
     ReferenceEntity,
+    ReverseConnectionEntity,
     URLEntity,
     ViewEntity,
     ViewEntityList,
-    ViewPropertyEntity,
 )
-from cognite.neat.rules.models.wrapped_entities import HasDataFilter, NodeTypeFilter, RawFilter
 
 from ._schema import DMSSchema
 
 if TYPE_CHECKING:
-    from cognite.neat.rules.models.information._rules import InformationRules
+    pass
 
 if sys.version_info >= (3, 11):
     pass
@@ -139,35 +140,6 @@ class DMSMetadata(BaseMetadata):
     def as_identifier(self) -> str:
         return repr(self.as_data_model_id())
 
-    @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
-
-    @classmethod
-    def from_data_model(cls, data_model: dm.DataModelApply, has_reference: bool) -> "DMSMetadata":
-        description, creator = cls._get_description_and_creator(data_model.description)
-        return cls(
-            schema_=SchemaCompleteness.complete,
-            data_model_type=DataModelType.solution if has_reference else DataModelType.enterprise,
-            space=data_model.space,
-            name=data_model.name or None,
-            description=description,
-            external_id=data_model.external_id,
-            version=data_model.version,
-            creator=creator,
-            created=datetime.now(),
-            updated=datetime.now(),
-        )
-
     def get_prefix(self) -> str:
         return self.space
 
@@ -177,8 +149,8 @@ class DMSProperty(SheetEntity):
     view_property: str = Field(alias="View Property")
     name: str | None = Field(alias="Name", default=None)
     description: str | None = Field(alias="Description", default=None)
-    connection: Literal["direct", "edge", "reverse"] | None = Field(None, alias="Connection")
-    value_type: DataType | ViewPropertyEntity | ViewEntity | DMSUnknownEntity = Field(alias="Value Type")
+    connection: Literal["direct"] | ReverseConnectionEntity | EdgeEntity | None = Field(None, alias="Connection")
+    value_type: DataType | ViewEntity | DMSUnknownEntity = Field(alias="Value Type")
     nullable: bool | None = Field(default=None, alias="Nullable")
     immutable: bool | None = Field(default=None, alias="Immutable")
     is_list: bool | None = Field(default=None, alias="Is List")
@@ -199,25 +171,23 @@ class DMSProperty(SheetEntity):
 
     @field_validator("value_type", mode="after")
     def connections_value_type(
-        cls, value: ViewPropertyEntity | ViewEntity | DMSUnknownEntity, info: ValidationInfo
-    ) -> DataType | ViewPropertyEntity | ViewEntity | DMSUnknownEntity:
+        cls, value: EdgeEntity | ViewEntity | DMSUnknownEntity, info: ValidationInfo
+    ) -> DataType | EdgeEntity | ViewEntity | DMSUnknownEntity:
         if (connection := info.data.get("connection")) is None:
             return value
         if connection == "direct" and not isinstance(value, ViewEntity | DMSUnknownEntity):
             raise ValueError(f"Direct relation must have a value type that points to a view, got {value}")
-        elif connection == "edge" and not isinstance(value, ViewEntity):
+        elif isinstance(connection, EdgeEntity) and not isinstance(value, ViewEntity):
             raise ValueError(f"Edge connection must have a value type that points to a view, got {value}")
-        elif connection == "reverse" and not isinstance(value, ViewPropertyEntity | ViewEntity):
-            raise ValueError(
-                f"Reverse connection must have a value type that points to a view or view property, got {value}"
-            )
+        elif isinstance(connection, ReverseConnectionEntity) and not isinstance(value, ViewEntity):
+            raise ValueError(f"Reverse connection must have a value type that points to a view, got {value}")
         return value
 
     @field_serializer("value_type", when_used="always")
     @staticmethod
-    def as_dms_type(value_type: DataType | ViewPropertyEntity | ViewEntity) -> str:
+    def as_dms_type(value_type: DataType | EdgeEntity | ViewEntity) -> str:
         if isinstance(value_type, DataType):
-            return value_type.dms._type
+            return value_type._suffix_extra_args(value_type.dms._type)
         else:
             return str(value_type)
 
@@ -228,6 +198,7 @@ class DMSContainer(SheetEntity):
     description: str | None = Field(alias="Description", default=None)
     reference: URLEntity | ReferenceEntity | None = Field(alias="Reference", default=None, union_mode="left_to_right")
     constraint: ContainerEntityList | None = Field(None, alias="Constraint")
+    used_for: Literal["node", "edge", "all"] | None = Field("all", alias="Used For")
     class_: ClassEntity = Field(alias="Class (linage)")
 
     def as_container(self) -> dm.ContainerApply:
@@ -244,22 +215,7 @@ class DMSContainer(SheetEntity):
             description=self.description,
             constraints=constraints or None,
             properties={},
-        )
-
-    @classmethod
-    def from_container(cls, container: dm.ContainerApply) -> "DMSContainer":
-        constraints: list[ContainerEntity] = []
-        for _, constraint_obj in (container.constraints or {}).items():
-            if isinstance(constraint_obj, dm.RequiresConstraint):
-                constraints.append(ContainerEntity.from_id(constraint_obj.require))
-            # UniquenessConstraint it handled in the properties
-        container_entity = ContainerEntity.from_id(container.as_id())
-        return cls(
-            class_=container_entity.as_class(),
-            container=container_entity,
-            name=container.name or None,
-            description=container.description,
-            constraint=constraints or None,
+            used_for=self.used_for,
         )
 
 
@@ -292,19 +248,27 @@ class DMSView(SheetEntity):
             properties={},
         )
 
-    @classmethod
-    def from_view(cls, view: dm.ViewApply, in_model: bool) -> "DMSView":
-        view_entity = ViewEntity.from_id(view.as_id())
-        class_entity = view_entity.as_class(skip_version=True)
-
-        return cls(
-            class_=class_entity,
-            view=view_entity,
-            description=view.description,
-            name=view.name,
-            implements=[ViewEntity.from_id(parent, _DEFAULT_VERSION) for parent in view.implements] or None,
-            in_model=in_model,
-        )
+
+class DMSNode(SheetEntity):
+    node: DMSNodeEntity = Field(alias="Node")
+    usage: Literal["type", "collection"] = Field(alias="Usage")
+    name: str | None = Field(alias="Name", default=None)
+    description: str | None = Field(alias="Description", default=None)
+
+    def as_node(self) -> dm.NodeApply:
+        if self.usage == "type":
+            return dm.NodeApply(space=self.node.space, external_id=self.node.external_id)
+        elif self.usage == "collection":
+            raise NotImplementedError("Collection nodes are not supported yet")
+        else:
+            raise ValueError(f"Unknown usage {self.usage}")
+
+
+class DMSEnum(SheetEntity):
+    collection: ClassEntity = Field(alias="Collection")
+    value: str = Field(alias="Value")
+    name: str | None = Field(alias="Name", default=None)
+    description: str | None = Field(alias="Description", default=None)
 
 
 class DMSRules(BaseRules):
@@ -312,6 +276,8 @@ class DMSRules(BaseRules):
     properties: SheetList[DMSProperty] = Field(alias="Properties")
     views: SheetList[DMSView] = Field(alias="Views")
     containers: SheetList[DMSContainer] | None = Field(None, alias="Containers")
+    enum: SheetList[DMSEnum] | None = Field(None, alias="Enum")
+    nodes: SheetList[DMSNode] | None = Field(None, alias="Nodes")
     last: "DMSRules | None" = Field(None, alias="Last", description="The previous version of the data model")
     reference: "DMSRules | None" = Field(None, alias="Reference")
 
@@ -401,82 +367,3 @@ class DMSRules(BaseRules):
         from ._exporter import _DMSExporter
 
         return _DMSExporter(self, include_pipeline, instance_space).to_schema()
-
-    def as_information_rules(self) -> "InformationRules":
-        from ._converter import _DMSRulesConverter
-
-        return _DMSRulesConverter(self).as_information_rules()
-
-    def as_domain_expert_rules(self) -> DomainRules:
-        from ._converter import _DMSRulesConverter
-
-        return _DMSRulesConverter(self).as_domain_rules()
-
-    def create_reference(
-        self, reference: "DMSRules", view_extension_mapping: dict[str, str], default_extension: str | None = None
-    ) -> None:
-        """Takes this data models and makes it into an extension of the reference data model.
-
-        The argument view_extension_mapping is a dictionary where the keys are views of this data model,
-        and each value is the view of the reference data model that the view should extend. For example:
-
-        ```python
-        view_extension_mapping = {"Pump": "Asset"}
-        ```
-
-        This would make the view "Pump" in this data model extend the view "Asset" in the reference data model.
-        Note that all the keys in the dictionary must be external ids of views in this data model,
-        and all the values must be external ids of views in the reference data model.
-
-        Args:
-            reference: The reference data model
-            view_extension_mapping: A dictionary mapping views in this data model to views in the reference data model
-            default_extension: The default view in the reference data model that views in this
-                data model should extend if no mapping is provided.
-
-        """
-        if self.reference is not None:
-            raise ValueError("Reference already exists")
-        self.reference = reference
-        view_by_external_id = {view.view.external_id: view for view in self.views}
-        ref_view_by_external_id = {view.view.external_id: view for view in reference.views}
-
-        if invalid_views := set(view_extension_mapping.keys()) - set(view_by_external_id.keys()):
-            raise ValueError(f"Views are not in this dat model {invalid_views}")
-        if invalid_views := set(view_extension_mapping.values()) - set(ref_view_by_external_id.keys()):
-            raise ValueError(f"Views are not in the reference data model {invalid_views}")
-        if default_extension and default_extension not in ref_view_by_external_id:
-            raise ValueError(f"Default extension view not in the reference data model {default_extension}")
-
-        properties_by_view_external_id: dict[str, dict[str, DMSProperty]] = defaultdict(dict)
-        for prop in self.properties:
-            properties_by_view_external_id[prop.view.external_id][prop.view_property] = prop
-
-        ref_properties_by_view_external_id: dict[str, dict[str, DMSProperty]] = defaultdict(dict)
-        for prop in reference.properties:
-            ref_properties_by_view_external_id[prop.view.external_id][prop.view_property] = prop
-
-        for view_external_id, view in view_by_external_id.items():
-            if view_external_id in view_extension_mapping:
-                ref_external_id = view_extension_mapping[view_external_id]
-            elif default_extension:
-                ref_external_id = default_extension
-            else:
-                continue
-
-            ref_view = ref_view_by_external_id[ref_external_id]
-            shared_properties = set(properties_by_view_external_id[view_external_id].keys()) & set(
-                ref_properties_by_view_external_id[ref_external_id].keys()
-            )
-            if shared_properties:
-                if view.implements is None:
-                    view.implements = [ref_view.view]
-                elif isinstance(view.implements, list) and ref_view.view not in view.implements:
-                    view.implements.append(ref_view.view)
-            for prop_name in shared_properties:
-                prop = properties_by_view_external_id[view_external_id][prop_name]
-                ref_prop = ref_properties_by_view_external_id[ref_external_id][prop_name]
-                if ref_prop.container and ref_prop.container_property:
-                    prop.container = ref_prop.container
-                    prop.container_property = ref_prop.container_property
-                prop.reference = ReferenceEntity.from_entity(ref_prop.view, ref_prop.view_property)