cognite-neat 0.77.3__py3-none-any.whl → 0.77.5__py3-none-any.whl

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

@@ -1,5 +1,6 @@
 import warnings
 from collections import defaultdict
+from collections.abc import Sequence
 from typing import Any, cast
 
 from cognite.client.data_classes import data_modeling as dm
@@ -11,7 +12,7 @@ from cognite.client.data_classes.data_modeling.views import (
 )
 
 from cognite.neat.rules import issues
-from cognite.neat.rules.models._base import DataModelType
+from cognite.neat.rules.models._base import DataModelType, ExtensionCategory, SchemaCompleteness
 from cognite.neat.rules.models.data_types import DataType
 from cognite.neat.rules.models.entities import (
     ContainerEntity,
@@ -58,19 +59,83 @@ class _DMSExporter:
         else:
             self._ref_views_by_id = {}
 
+        self.is_addition = (
+            rules.metadata.schema_ is SchemaCompleteness.extended
+            and rules.metadata.extension is ExtensionCategory.addition
+        )
+        self.is_reshape = (
+            rules.metadata.schema_ is SchemaCompleteness.extended
+            and rules.metadata.extension is ExtensionCategory.reshape
+        )
+        self.is_rebuild = (
+            rules.metadata.schema_ is SchemaCompleteness.extended
+            and rules.metadata.extension is ExtensionCategory.rebuild
+        )
+
     def to_schema(self) -> DMSSchema:
         rules = self.rules
-        container_properties_by_id, view_properties_by_id = self._gather_properties()
+        container_properties_by_id, view_properties_by_id = self._gather_properties(list(self.rules.properties))
+
+        # If we are reshaping or rebuilding, and there are no properties in the current rules, we will
+        # include those properties from the last rules.
+        if rules.last and (self.is_reshape or self.is_rebuild):
+            selected_views = {view.view for view in rules.views}
+            selected_properties = [
+                prop
+                for prop in rules.last.properties
+                if prop.view in selected_views and prop.view.as_id() not in view_properties_by_id
+            ]
+            self._update_with_properties(
+                selected_properties, container_properties_by_id, view_properties_by_id, include_new_containers=True
+            )
+
+        # We need to include the properties from the last rules as well to create complete containers and view
+        # depending on the type of extension.
+        if rules.last and self.is_addition:
+            selected_properties = [
+                prop for prop in rules.last.properties if (prop.view.as_id() in view_properties_by_id)
+            ]
+            self._update_with_properties(selected_properties, container_properties_by_id, view_properties_by_id)
+        elif rules.last and (self.is_reshape or self.is_rebuild):
+            selected_properties = [
+                prop
+                for prop in rules.last.properties
+                if prop.container and prop.container.as_id() in container_properties_by_id
+            ]
+            self._update_with_properties(selected_properties, container_properties_by_id, None)
+
         containers = self._create_containers(container_properties_by_id)
 
         views, node_types = self._create_views_with_node_types(view_properties_by_id)
 
+        last_schema: DMSSchema | None = None
+        if self.rules.last:
+            last_schema = self.rules.last.as_schema()
+            # Remove the views that are in the current model, last + current should equal the full model
+            # without any duplicates
+            last_schema.views = ViewApplyDict(
+                {view_id: view for view_id, view in last_schema.views.items() if view_id not in views}
+            )
+            last_schema.containers = ContainerApplyDict(
+                {
+                    container_id: container
+                    for container_id, container in last_schema.containers.items()
+                    if container_id not in containers
+                }
+            )
+
         views_not_in_model = {view.view.as_id() for view in rules.views if not view.in_model}
+        if rules.last and self.is_addition:
+            views_not_in_model.update({view.view.as_id() for view in rules.last.views if not view.in_model})
+
         data_model = rules.metadata.as_data_model()
-        data_model.views = sorted(
-            [view_id for view_id in views.keys() if view_id not in views_not_in_model],
-            key=lambda v: v.as_tuple(),  # type: ignore[union-attr]
-        )
+
+        data_model_views = [view_id for view_id in views if view_id not in views_not_in_model]
+        if last_schema and self.is_addition:
+            data_model_views.extend([view_id for view_id in last_schema.views if view_id not in views_not_in_model])
+
+        # Sorting to ensure deterministic order
+        data_model.views = sorted(data_model_views, key=lambda v: v.as_tuple())  # type: ignore[union-attr]
 
         spaces = self._create_spaces(rules.metadata, containers, views, data_model)
 
@@ -86,7 +151,8 @@ class _DMSExporter:
 
         if self._ref_schema:
             output.reference = self._ref_schema
-
+        if last_schema:
+            output.last = last_schema
         return output
 
     def _create_spaces(
@@ -112,8 +178,18 @@ class _DMSExporter:
         self,
         view_properties_by_id: dict[dm.ViewId, list[DMSProperty]],
     ) -> tuple[ViewApplyDict, NodeApplyDict]:
-        views = ViewApplyDict([dms_view.as_view() for dms_view in self.rules.views])
-        dms_view_by_id = {dms_view.view.as_id(): dms_view for dms_view in self.rules.views}
+        input_views = list(self.rules.views)
+        if self.rules.last:
+            existing = {view.view.as_id() for view in input_views}
+            modified_views = [
+                v
+                for v in self.rules.last.views
+                if v.view.as_id() in view_properties_by_id and v.view.as_id() not in existing
+            ]
+            input_views.extend(modified_views)
+
+        views = ViewApplyDict([dms_view.as_view() for dms_view in input_views])
+        dms_view_by_id = {dms_view.view.as_id(): dms_view for dms_view in input_views}
 
         for view_id, view in views.items():
             view.properties = {}
@@ -176,9 +252,17 @@ class _DMSExporter:
         self,
         container_properties_by_id: dict[dm.ContainerId, list[DMSProperty]],
     ) -> ContainerApplyDict:
-        containers = dm.ContainerApplyList(
-            [dms_container.as_container() for dms_container in self.rules.containers or []]
-        )
+        containers = list(self.rules.containers or [])
+        if self.rules.last:
+            existing = {container.container.as_id() for container in containers}
+            modified_containers = [
+                c
+                for c in self.rules.last.containers or []
+                if c.container.as_id() in container_properties_by_id and c.container.as_id() not in existing
+            ]
+            containers.extend(modified_containers)
+
+        containers = dm.ContainerApplyList([dms_container.as_container() for dms_container in containers])
         container_to_drop = set()
         for container in containers:
             container_id = container.as_id()
@@ -232,10 +316,13 @@ class _DMSExporter:
                 }
         return ContainerApplyDict([container for container in containers if container.as_id() not in container_to_drop])
 
-    def _gather_properties(self) -> tuple[dict[dm.ContainerId, list[DMSProperty]], dict[dm.ViewId, list[DMSProperty]]]:
+    @staticmethod
+    def _gather_properties(
+        properties: Sequence[DMSProperty],
+    ) -> tuple[dict[dm.ContainerId, list[DMSProperty]], dict[dm.ViewId, list[DMSProperty]]]:
         container_properties_by_id: dict[dm.ContainerId, list[DMSProperty]] = defaultdict(list)
         view_properties_by_id: dict[dm.ViewId, list[DMSProperty]] = defaultdict(list)
-        for prop in self.rules.properties:
+        for prop in properties:
             view_id = prop.view.as_id()
             view_properties_by_id[view_id].append(prop)
 
@@ -245,6 +332,32 @@ class _DMSExporter:
 
         return container_properties_by_id, view_properties_by_id
 
+    @classmethod
+    def _update_with_properties(
+        cls,
+        selected_properties: Sequence[DMSProperty],
+        container_properties_by_id: dict[dm.ContainerId, list[DMSProperty]],
+        view_properties_by_id: dict[dm.ViewId, list[DMSProperty]] | None,
+        include_new_containers: bool = False,
+    ) -> None:
+        view_properties_by_id = view_properties_by_id or {}
+        last_container_properties_by_id, last_view_properties_by_id = cls._gather_properties(selected_properties)
+
+        for container_id, properties in last_container_properties_by_id.items():
+            # Only add the container properties that are not already present, and do not overwrite.
+            if (container_id in container_properties_by_id) or include_new_containers:
+                existing = {prop.container_property for prop in container_properties_by_id.get(container_id, [])}
+                container_properties_by_id[container_id].extend(
+                    [prop for prop in properties if prop.container_property not in existing]
+                )
+
+        if view_properties_by_id:
+            for view_id, properties in last_view_properties_by_id.items():
+                existing = {prop.view_property for prop in view_properties_by_id[view_id]}
+                view_properties_by_id[view_id].extend(
+                    [prop for prop in properties if prop.view_property not in existing]
+                )
+
     def _create_view_filter(
         self,
         view: dm.ViewApply,
@@ -290,8 +403,9 @@ class _DMSExporter:
             # HasData or not provided (this is the default)
             return HasDataFilter(inner=[ContainerEntity.from_id(id_) for id_ in ref_containers])
 
+    @classmethod
     def _create_view_property(
-        self, prop: DMSProperty, view_properties_by_id: dict[dm.ViewId, list[DMSProperty]]
+        cls, prop: DMSProperty, view_properties_by_id: dict[dm.ViewId, list[DMSProperty]]
     ) -> ViewPropertyApply | None:
         if prop.container and prop.container_property:
             container_prop_identifier = prop.container_property
@@ -335,7 +449,7 @@ class _DMSExporter:
                 edge_cls = SingleEdgeConnectionApply
 
             return edge_cls(
-                type=self._create_edge_type_from_prop(prop),
+                type=cls._create_edge_type_from_prop(prop),
                 source=source_view_id,
                 direction="outwards",
                 name=prop.name,
@@ -376,7 +490,7 @@ class _DMSExporter:
                     dm.MultiEdgeConnectionApply if prop.is_list in [True, None] else SingleEdgeConnectionApply
                 )
                 return inwards_edge_cls(
-                    type=self._create_edge_type_from_prop(reverse_prop or prop),
+                    type=cls._create_edge_type_from_prop(reverse_prop or prop),
                     source=source_view_id,
                     name=prop.name,
                     description=prop.description,

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

@@ -2,13 +2,13 @@ import math
 import re
 import sys
 import warnings
-from collections.abc import Callable
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Literal
 
 from cognite.client import data_modeling as dm
-from pydantic import Field, field_serializer, field_validator, model_serializer, model_validator
-from pydantic_core.core_schema import SerializationInfo, ValidationInfo
+from pydantic import Field, field_serializer, field_validator, model_validator
+from pydantic.main import IncEx
+from pydantic_core.core_schema import ValidationInfo
 
 from cognite.neat.rules import issues
 from cognite.neat.rules.issues import MultiValueError
@@ -49,16 +49,16 @@ if TYPE_CHECKING:
     from cognite.neat.rules.models.information._rules import InformationRules
 
 if sys.version_info >= (3, 11):
-    from typing import Self
+    pass
 else:
-    from typing_extensions import Self
+    pass
 
 _DEFAULT_VERSION = "1"
 
 
 class DMSMetadata(BaseMetadata):
     role: ClassVar[RoleTypes] = RoleTypes.dms_architect
-    data_model_type: DataModelType = Field(DataModelType.solution, alias="dataModelType")
+    data_model_type: DataModelType = Field(DataModelType.enterprise, alias="dataModelType")
     schema_: SchemaCompleteness = Field(alias="schema")
     extension: ExtensionCategory = ExtensionCategory.addition
     space: ExternalIdType
@@ -146,10 +146,11 @@ class DMSMetadata(BaseMetadata):
         return description, creator
 
     @classmethod
-    def from_data_model(cls, data_model: dm.DataModelApply) -> "DMSMetadata":
+    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,
@@ -263,13 +264,20 @@ class DMSView(SheetEntity):
 
     def as_view(self) -> dm.ViewApply:
         view_id = self.view.as_id()
+        implements = [parent.as_id() for parent in self.implements or []] or None
+        if implements is None and isinstance(self.reference, ReferenceEntity):
+            # Fallback to the reference if no implements are provided
+            parent = self.reference.as_view_id()
+            if (parent.space, parent.external_id) != (view_id.space, view_id.external_id):
+                implements = [parent]
+
         return dm.ViewApply(
             space=view_id.space,
             external_id=view_id.external_id,
             version=view_id.version or _DEFAULT_VERSION,
             name=self.name or None,
             description=self.description,
-            implements=[parent.as_id() for parent in self.implements or []] or None,
+            implements=implements,
             properties={},
         )
 
@@ -333,17 +341,37 @@ class DMSRules(BaseRules):
             raise MultiValueError([error for error in issue_list if isinstance(error, issues.NeatValidationError)])
         return self
 
-    @model_serializer(mode="wrap", when_used="always")
-    def dms_rules_serialization(
+    def dump(
         self,
-        handler: Callable,
-        info: SerializationInfo,
+        mode: Literal["python", "json"] = "python",
+        by_alias: bool = False,
+        exclude: IncEx = None,
+        exclude_none: bool = False,
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        as_reference: bool = False,
     ) -> dict[str, Any]:
         from ._serializer import _DMSRulesSerializer
 
-        dumped = cast(dict[str, Any], handler(self, info))
+        dumped = self.model_dump(
+            mode=mode,
+            by_alias=by_alias,
+            exclude=exclude,
+            exclude_none=exclude_none,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+        )
         space, version = self.metadata.space, self.metadata.version
-        return _DMSRulesSerializer(info, space, version).clean(dumped)
+        serializer = _DMSRulesSerializer(by_alias, space, version)
+        clean = serializer.clean(dumped, as_reference)
+        last = "Last" if by_alias else "last"
+        if last_dump := clean.get(last):
+            clean[last] = serializer.clean(last_dump, False)
+        reference = "Reference" if by_alias else "reference"
+        if self.reference and (ref_dump := clean.get(reference)):
+            space, version = self.reference.metadata.space, self.reference.metadata.version
+            clean[reference] = _DMSRulesSerializer(by_alias, space, version).clean(ref_dump, True)
+        return clean
 
     def as_schema(self, include_pipeline: bool = False, instance_space: str | None = None) -> DMSSchema:
         from ._exporter import _DMSExporter
@@ -359,19 +387,3 @@ class DMSRules(BaseRules):
         from ._converter import _DMSRulesConverter
 
         return _DMSRulesConverter(self).as_domain_rules()
-
-    def reference_self(self) -> Self:
-        new_rules = self.model_copy(deep=True)
-        for prop in new_rules.properties:
-            prop.reference = ReferenceEntity(
-                prefix=prop.view.prefix, suffix=prop.view.suffix, version=prop.view.version, property=prop.property_
-            )
-        view: DMSView
-        for view in new_rules.views:
-            view.reference = ReferenceEntity(
-                prefix=view.view.prefix, suffix=view.view.suffix, version=view.view.version
-            )
-        container: DMSContainer
-        for container in new_rules.containers or []:
-            container.reference = ReferenceEntity(prefix=container.container.prefix, suffix=container.container.suffix)
-        return new_rules

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

@@ -3,9 +3,7 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, Literal, cast, overload
 
-from pydantic import BaseModel
-
-from cognite.neat.rules.models._base import DataModelType, ExtensionCategory, SchemaCompleteness
+from cognite.neat.rules.models._base import DataModelType, ExtensionCategory, SchemaCompleteness, _add_alias
 from cognite.neat.rules.models.data_types import DataType
 from cognite.neat.rules.models.entities import (
     ClassEntity,
@@ -44,8 +42,9 @@ class DMSMetadataInput:
             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]
-            extension=data.get("extension", "addition"),
-            data_model_type=data.get("data_model_type", "solution"),
+            # 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"),
@@ -355,9 +354,3 @@ class DMSRulesInput:
             Last=last,
             Reference=reference,
         )
-
-
-def _add_alias(data: dict[str, Any], base_model: type[BaseModel]) -> None:
-    for field_name, field_ in base_model.model_fields.items():
-        if field_name not in data and field_.alias in data:
-            data[field_name] = data[field_.alias]

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

@@ -516,10 +516,11 @@ class DMSSchema:
         defined_spaces = self.spaces.copy()
         defined_containers = self.containers.copy()
         defined_views = self.views.copy()
-        if self.reference:
-            defined_spaces |= self.reference.spaces
-            defined_containers |= self.reference.containers
-            defined_views |= self.reference.views
+        for other_schema in [self.reference, self.last]:
+            if other_schema:
+                defined_spaces |= other_schema.spaces
+                defined_containers |= other_schema.containers
+                defined_views |= other_schema.views
 
         for container in self.containers.values():
             if container.space not in defined_spaces:

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

@@ -1,9 +1,8 @@
 from typing import Any, ClassVar, cast
 
-from pydantic_core.core_schema import SerializationInfo
-
 from cognite.neat.rules.models import DMSRules
 from cognite.neat.rules.models.dms import DMSContainer, DMSProperty, DMSView
+from cognite.neat.rules.models.entities import ReferenceEntity, ViewEntity
 
 
 class _DMSRulesSerializer:
@@ -12,7 +11,7 @@ class _DMSRulesSerializer:
     VIEWS_FIELDS: ClassVar[list[str]] = ["class_", "view", "implements"]
     CONTAINERS_FIELDS: ClassVar[list[str]] = ["class_", "container"]
 
-    def __init__(self, info: SerializationInfo, default_space: str, default_version: str) -> None:
+    def __init__(self, by_alias: bool, default_space: str, default_version: str) -> None:
         self.default_space = f"{default_space}:"
         self.default_version = f"version={default_version}"
         self.default_version_wrapped = f"({self.default_version})"
@@ -25,22 +24,25 @@ class _DMSRulesSerializer:
         self.container_name = "containers"
         self.metadata_name = "metadata"
         self.prop_view = "view"
+        self.prop_container = "container"
         self.prop_view_property = "view_property"
         self.prop_value_type = "value_type"
         self.view_view = "view"
         self.view_implements = "implements"
         self.container_container = "container"
         self.container_constraint = "constraint"
+        self.reference = "Reference" if by_alias else "reference"
 
-        if info.by_alias:
+        if by_alias:
             self.properties_fields = [
                 DMSProperty.model_fields[field].alias or field for field in self.properties_fields
             ]
             self.views_fields = [DMSView.model_fields[field].alias or field for field in self.views_fields]
-            self.container_fields = [
+            self.containers_fields = [
                 DMSContainer.model_fields[field].alias or field for field in self.containers_fields
             ]
             self.prop_view = DMSProperty.model_fields[self.prop_view].alias or self.prop_view
+            self.prop_container = DMSProperty.model_fields[self.prop_container].alias or self.prop_container
             self.prop_view_property = DMSProperty.model_fields[self.prop_view_property].alias or self.prop_view_property
             self.prop_value_type = DMSProperty.model_fields[self.prop_value_type].alias or self.prop_value_type
             self.view_view = DMSView.model_fields[self.view_view].alias or self.view_view
@@ -56,22 +58,7 @@ class _DMSRulesSerializer:
             self.container_name = DMSRules.model_fields[self.container_name].alias or self.container_name
             self.metadata_name = DMSRules.model_fields[self.metadata_name].alias or self.metadata_name
 
-        if isinstance(info.exclude, dict):
-            # Just for happy mypy
-            exclude = cast(dict, info.exclude)
-            self.exclude_properties = exclude.get("properties", {}).get("__all__", set())
-            self.exclude_views = exclude.get("views", {}).get("__all__", set()) or set()
-            self.exclude_containers = exclude.get("containers", {}).get("__all__", set()) or set()
-            self.metadata_exclude = exclude.get("metadata", set()) or set()
-            self.exclude_top = {k for k, v in exclude.items() if not v}
-        else:
-            self.exclude_top = set(info.exclude or {})
-            self.exclude_properties = set()
-            self.exclude_views = set()
-            self.exclude_containers = set()
-            self.metadata_exclude = set()
-
-    def clean(self, dumped: dict[str, Any]) -> dict[str, Any]:
+    def clean(self, dumped: dict[str, Any], as_reference: bool) -> dict[str, Any]:
         # Sorting to get a deterministic order
         dumped[self.prop_name] = sorted(
             dumped[self.prop_name]["data"], key=lambda p: (p[self.prop_view], p[self.prop_view_property])
@@ -83,16 +70,29 @@ class _DMSRulesSerializer:
             dumped.pop(self.container_name, None)
 
         for prop in dumped[self.prop_name]:
+            if as_reference:
+                view_entity = cast(ViewEntity, ViewEntity.load(prop[self.prop_view]))
+                prop[self.reference] = str(
+                    ReferenceEntity(
+                        prefix=view_entity.prefix,
+                        suffix=view_entity.suffix,
+                        version=view_entity.version,
+                        property=prop[self.prop_view_property],
+                    )
+                )
             for field_name in self.properties_fields:
+                if as_reference and field_name == self.prop_container:
+                    # When dumping as reference, the container should keep the default space for easy copying
+                    # over to user sheets.
+                    continue
                 if value := prop.get(field_name):
                     prop[field_name] = value.removeprefix(self.default_space).removesuffix(self.default_version_wrapped)
             # Value type can have a property as well
             prop[self.prop_value_type] = prop[self.prop_value_type].replace(self.default_version, "")
-            if self.exclude_properties:
-                for field in self.exclude_properties:
-                    prop.pop(field, None)
 
         for view in dumped[self.view_name]:
+            if as_reference:
+                view[self.reference] = view[self.view_view]
             for field_name in self.views_fields:
                 if value := view.get(field_name):
                     view[field_name] = value.removeprefix(self.default_space).removesuffix(self.default_version_wrapped)
@@ -101,11 +101,10 @@ class _DMSRulesSerializer:
                     parent.strip().removeprefix(self.default_space).removesuffix(self.default_version_wrapped)
                     for parent in value.split(",")
                 )
-            if self.exclude_views:
-                for field in self.exclude_views:
-                    view.pop(field, None)
 
         for container in dumped.get(self.container_name, []):
+            if as_reference:
+                container[self.reference] = container[self.container_container]
             for field_name in self.containers_fields:
                 if value := container.get(field_name):
                     container[field_name] = value.removeprefix(self.default_space)
@@ -114,13 +113,4 @@ class _DMSRulesSerializer:
                 container[self.container_constraint] = ",".join(
                     constraint.strip().removeprefix(self.default_space) for constraint in value.split(",")
                 )
-            if self.exclude_containers:
-                for field in self.exclude_containers:
-                    container.pop(field, None)
-
-        if self.metadata_exclude:
-            for field in self.metadata_exclude:
-                dumped[self.metadata_name].pop(field, None)
-        for field in self.exclude_top:
-            dumped.pop(field, None)
         return dumped