[NEAT-230] 🏗Extending solution model (#460)

* tests: clearer names

* fix: bug in DMSToRules step

* fix: DMSImporter download enterprise model

* fix; conversion dms to info

* fix: added missing conversion

* fix: logic bug

* refactor: Added Svein Harald spredsheet

* fix: Include metadata sheet

* tests: extending excel importer to svein harald

* fix: bug in ExcelImporter

* refactor: first pass at refactoring DMS validation

* refactor: update validation

* refactor: updated olav to complete

* refactor: update validation

* refactor: update rules

* tests: updated failing test

* docs: Added Svein Harald DMS rules

* tests: Added Svein Harald to tests

* fix: Bug in DMS serializer

* tests: failing tests

* tests: for bad conversion

* fix: missing parent in InfoInput rules

* refactor; fix svein harald

* refactor; revert changes

* refactor; setup shell

* docs; downloading section

* tests: Added test for
exporting Olav updated dms model

* fix; bug in DMS importer step

* refactor; setup shell for new dump method

* refactor: update serialization

* tests: updated test

* tests: updated test

* refactor: update YAML exporter

* feat: implemented correct as reference for DMS

* tests: updated tests

* tests: updated

* fix; few issues

* refactor; implemented as reference for information model

* fix; DMS serializer

* refactor: first pass on Olav

* docs; starting on next section

* fix: Svein Harald

* test: created failing test for newly found bug

* fix: as schema with last

* fix: validation

* fix: validatin and as schema part 1

* feat: Smart conversion of information rules

* docs: finished rubild section and spreadsheet

* tests: Added failing test for Olav rebuild as schema

* refactor: fix

* fix; implementes fallback to reference

* fix: introduced bug

* docs; finished part 4 first draft

* tests: updated test

* tests: updated integration test

* docs; update links

* build; changelog

Makefile

@@ -1,6 +1,6 @@
 .PHONY: run-explorer run-tests run-linters build-ui build-python build-docker run-docker compose-up
 
-version="0.77.4"
+version="0.77.5"
 run-explorer:
 	@echo "Running explorer API server..."
 	# open "http://localhost:8000/static/index.html" || true

cognite/neat/_version.py

@@ -1 +1 @@
-__version__ = "0.77.4"
+__version__ = "0.77.5"

cognite/neat/rules/exporters/_rules2excel.py

@@ -116,17 +116,17 @@ def export(self, rules: Rules) -> Workbook:
             }
 
             if self.dump_as == "last":
-                dumped_last_rules = rules.model_dump(by_alias=True)
+                dumped_last_rules = rules.dump(by_alias=True)
                 if rules.reference:
-                    dumped_reference_rules = rules.reference.model_dump(by_alias=True)
+                    dumped_reference_rules = rules.reference.dump(by_alias=True, as_reference=True)
             elif self.dump_as == "reference":
-                dumped_reference_rules = rules.reference_self().model_dump(by_alias=True)
+                dumped_reference_rules = rules.dump(by_alias=True, as_reference=True)
         else:
-            dumped_user_rules = rules.model_dump(by_alias=True)
+            dumped_user_rules = rules.dump(by_alias=True)
             if rules.last:
-                dumped_last_rules = rules.last.model_dump(by_alias=True)
+                dumped_last_rules = rules.last.dump(by_alias=True)
             if rules.reference:
-                dumped_reference_rules = rules.reference.model_dump(by_alias=True)
+                dumped_reference_rules = rules.reference.dump(by_alias=True, as_reference=True)
 
         self._write_metadata_sheet(workbook, dumped_user_rules["Metadata"])
         self._write_sheets(workbook, dumped_user_rules, rules)

cognite/neat/rules/exporters/_rules2yaml.py

@@ -69,10 +69,10 @@ def export(self, rules: Rules) -> str:
         rules = self._convert_to_output_role(rules, self.output_role)
         # model_dump_json ensures that the output is in JSON format,
         # if we don't do this, we will get Enums and other types that are not serializable to YAML
-        json_output = rules.model_dump_json()
+        json_output = rules.dump(mode="json")
         if self.output == "json":
-            return json_output
+            return json.dumps(json_output)
         elif self.output == "yaml":
-            return yaml.safe_dump(json.loads(json_output))
+            return yaml.safe_dump(json_output)
         else:
             raise ValueError(f"Invalid output: {self.output}. Valid options are {self.format_option}")

cognite/neat/rules/models/_base.py

@@ -7,10 +7,9 @@
 import math
 import sys
 import types
-from abc import abstractmethod
 from collections.abc import Callable, Iterator
 from functools import wraps
-from typing import Annotated, Any, ClassVar, Generic, TypeAlias, TypeVar
+from typing import Annotated, Any, ClassVar, Generic, Literal, TypeAlias, TypeVar
 
 import pandas as pd
 from pydantic import (
@@ -26,13 +25,12 @@
     model_validator,
 )
 from pydantic.fields import FieldInfo
+from pydantic.main import IncEx
 
 if sys.version_info >= (3, 11):
     from enum import StrEnum
-    from typing import Self
 else:
     from backports.strenum import StrEnum
-    from typing_extensions import Self
 
 
 METADATA_VALUE_MAX_LENGTH = 5120
@@ -259,26 +257,33 @@ class BaseRules(RuleModel):
 
     Args:
         metadata: Data model metadata
-        classes: Classes defined in the data model
-        properties: Class properties defined in the data model with accompanying transformation rules
-                    to transform data from source to target representation
-        prefixes: Prefixes used in the data model. Defaults to PREFIXES
-        instances: Instances defined in the data model. Defaults to None
         validators_to_skip: List of validators to skip. Defaults to []
     """
 
     metadata: BaseMetadata
 
-    @abstractmethod
-    def reference_self(self) -> Self:
+    def dump(
+        self,
+        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]:
+        """Dump the model to a dictionary.
+
+        This is used in the Exporters to dump rules in the required format.
         """
-        Returns a copy of the rules with reference fields set to itself
-
-        For example, if the rules have a property with a reference field, then
-        the reference field will be set to the property itself. This is used when
-        exporting a reference model.
-        """
-        raise NotImplementedError
+        return self.model_dump(
+            mode=mode,
+            by_alias=by_alias,
+            exclude=exclude,
+            exclude_none=exclude_none,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+        )
 
 
 # An sheet entity is either a class or a property.

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.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 @@ def __init__(
         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,8 +151,8 @@ def to_schema(self) -> DMSSchema:
 
         if self._ref_schema:
             output.reference = self._ref_schema
-        if self.rules.last:
-            output.last = self.rules.last.as_schema()
+        if last_schema:
+            output.last = last_schema
         return output
 
     def _create_spaces(
@@ -113,8 +178,18 @@ def _create_views_with_node_types(
         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 = {}
@@ -177,9 +252,17 @@ def _create_containers(
         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()
@@ -233,10 +316,13 @@ def _create_containers(
                 }
         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)
 
@@ -246,6 +332,32 @@ def _gather_properties(self) -> tuple[dict[dm.ContainerId, list[DMSProperty]], d
 
         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,
@@ -291,8 +403,9 @@ def _create_view_filter(
             # 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
@@ -336,7 +449,7 @@ def _create_view_property(
                 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,
@@ -377,7 +490,7 @@ def _create_view_property(
                     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,