cognite-neat 0.98.0__py3-none-any.whl → 0.99.1__py3-none-any.whl

cognite/neat/_rules/transformers/_mapping.py

@@ -1,13 +1,16 @@
+import warnings
 from abc import ABC
 from collections import defaultdict
+from functools import cached_property
+from typing import Any, ClassVar, Literal
 
+from cognite.neat._issues.errors import NeatValueError
+from cognite.neat._issues.warnings import NeatValueWarning, PropertyOverwritingWarning
 from cognite.neat._rules._shared import JustRules, OutRules
-from cognite.neat._rules.models import DMSRules, InformationRules
-from cognite.neat._rules.models._base_rules import ClassRef
-from cognite.neat._rules.models.dms import DMSProperty
-from cognite.neat._rules.models.entities import ClassEntity
-from cognite.neat._rules.models.information import InformationClass
-from cognite.neat._rules.models.mapping import RuleMapping
+from cognite.neat._rules.models import DMSRules, SheetList
+from cognite.neat._rules.models.data_types import Enum
+from cognite.neat._rules.models.dms import DMSEnum, DMSProperty, DMSView
+from cognite.neat._rules.models.entities import ViewEntity
 
 from ._base import RulesTransformer
 
@@ -96,7 +99,7 @@ class MapOneToOne(MapOntoTransformers):
         return JustRules(solution)
 
 
-class RuleMapper(RulesTransformer[InformationRules, InformationRules]):
+class RuleMapper(RulesTransformer[DMSRules, DMSRules]):
     """Maps properties and classes using the given mapping.
 
     **Note**: This transformer mutates the input rules.
@@ -106,30 +109,102 @@ class RuleMapper(RulesTransformer[InformationRules, InformationRules]):
 
     """
 
-    def __init__(self, mapping: RuleMapping) -> None:
+    _mapping_fields: ClassVar[frozenset[str]] = frozenset(
+        ["connection", "value_type", "nullable", "immutable", "is_list", "default", "index", "constraint"]
+    )
+
+    def __init__(self, mapping: DMSRules, data_type_conflict: Literal["overwrite"] = "overwrite") -> None:
         self.mapping = mapping
+        self.data_type_conflict = data_type_conflict
+
+    @cached_property
+    def _view_by_entity_id(self) -> dict[str, DMSView]:
+        return {view.view.external_id: view for view in self.mapping.views}
 
-    def transform(self, rules: InformationRules | OutRules[InformationRules]) -> JustRules[InformationRules]:
+    @cached_property
+    def _property_by_view_property(self) -> dict[tuple[str, str], DMSProperty]:
+        return {(prop.view.external_id, prop.view_property): prop for prop in self.mapping.properties}
+
+    def transform(self, rules: DMSRules | OutRules[DMSRules]) -> JustRules[DMSRules]:
+        if self.data_type_conflict != "overwrite":
+            raise NeatValueError(f"Invalid data_type_conflict: {self.data_type_conflict}")
         input_rules = self._to_rules(rules)
+        new_rules = input_rules.model_copy(deep=True)
+        new_rules.metadata.version += "_mapped"
 
-        destination_by_source = self.mapping.properties.as_destination_by_source()
-        destination_cls_by_source = self.mapping.classes.as_destination_by_source()
-        used_destination_classes: set[ClassEntity] = set()
-        for prop in input_rules.properties:
-            if destination_prop := destination_by_source.get(prop.as_reference()):
-                prop.class_ = destination_prop.class_
-                prop.property_ = destination_prop.property_
-                used_destination_classes.add(destination_prop.class_)
-            elif destination_cls := destination_cls_by_source.get(ClassRef(Class=prop.class_)):
-                # If the property is not in the mapping, but the class is,
-                # then we should map the class to the destination
-                prop.class_ = destination_cls.class_
-
-        for cls_ in input_rules.classes:
-            if destination_cls := destination_cls_by_source.get(cls_.as_reference()):
-                cls_.class_ = destination_cls.class_
-        existing_classes = {cls_.class_ for cls_ in input_rules.classes}
-        for new_cls in used_destination_classes - existing_classes:
-            input_rules.classes.append(InformationClass(class_=new_cls))
-
-        return JustRules(input_rules)
+        for view in new_rules.views:
+            if mapping_view := self._view_by_entity_id.get(view.view.external_id):
+                view.implements = mapping_view.implements
+
+        for prop in new_rules.properties:
+            key = (prop.view.external_id, prop.view_property)
+            if key not in self._property_by_view_property:
+                continue
+            mapping_prop = self._property_by_view_property[key]
+            to_overwrite, conflicts = self._find_overwrites(prop, mapping_prop)
+            if conflicts and self.data_type_conflict == "overwrite":
+                warnings.warn(
+                    PropertyOverwritingWarning(prop.view.as_id(), "view", prop.view_property, tuple(conflicts)),
+                    stacklevel=2,
+                )
+            elif conflicts:
+                raise NeatValueError(f"Conflicting properties for {prop.view}.{prop.view_property}: {conflicts}")
+            for field_name, value in to_overwrite.items():
+                setattr(prop, field_name, value)
+            prop.container = mapping_prop.container
+            prop.container_property = mapping_prop.container_property
+
+        # Add missing views used as value types
+        existing_views = {view.view for view in new_rules.views}
+        new_value_types = {
+            prop.value_type
+            for prop in new_rules.properties
+            if isinstance(prop.value_type, ViewEntity) and prop.value_type not in existing_views
+        }
+        for new_value_type in new_value_types:
+            if mapping_view := self._view_by_entity_id.get(new_value_type.external_id):
+                new_rules.views.append(mapping_view)
+            else:
+                warnings.warn(NeatValueWarning(f"View {new_value_type} not found in mapping"), stacklevel=2)
+
+        # Add missing enums
+        existing_enum_collections = {item.collection for item in new_rules.enum or []}
+        new_enums = {
+            prop.value_type.collection
+            for prop in new_rules.properties
+            if isinstance(prop.value_type, Enum) and prop.value_type.collection not in existing_enum_collections
+        }
+        if new_enums:
+            new_rules.enum = new_rules.enum or SheetList[DMSEnum]([])
+            for item in self.mapping.enum or []:
+                if item.collection in new_enums:
+                    new_rules.enum.append(item)
+
+        return JustRules(new_rules)
+
+    def _find_overwrites(self, prop: DMSProperty, mapping_prop: DMSProperty) -> tuple[dict[str, Any], list[str]]:
+        """Finds the properties that need to be overwritten and returns them.
+
+        In addition, conflicting properties are returned. Note that overwriting properties that are
+        originally None is not considered a conflict. Thus, you can have properties to overwrite but no
+        conflicts.
+
+        Args:
+            prop: The property to compare.
+            mapping_prop: The property to compare against.
+
+        Returns:
+            A tuple with the properties to overwrite and the conflicting properties.
+
+        """
+        to_overwrite: dict[str, Any] = {}
+        conflicts: list[str] = []
+        for field_name in self._mapping_fields:
+            mapping_value = getattr(mapping_prop, field_name)
+            source_value = getattr(prop, field_name)
+            if mapping_value != source_value:
+                to_overwrite[field_name] = mapping_value
+                if source_value is not None:
+                    # These are used for warnings so we use the alias to make it more readable for the user
+                    conflicts.append(mapping_prop.model_fields[field_name].alias or field_name)
+        return to_overwrite, conflicts

cognite/neat/_rules/transformers/_pipelines.py

@@ -23,7 +23,7 @@ class ImporterPipeline(RulesPipeline[InputRules, VerifiedRules]):
 
     @classmethod
     def _create_pipeline(cls, importer: BaseImporter[InputRules], role: RoleTypes | None = None) -> "ImporterPipeline":
-        items: list[RulesTransformer] = [VerifyAnyRules(errors="continue")]
+        items: list[RulesTransformer] = [VerifyAnyRules(errors="continue", validate=True)]
         if role is not None:
             out_cls = VERIFIED_RULES_BY_ROLE[role]
             items.append(ConvertToRules(out_cls))

cognite/neat/_rules/transformers/_verification.py

@@ -1,7 +1,7 @@
 from abc import ABC
 from typing import Any, Literal
 
-from cognite.neat._issues import IssueList, NeatError, NeatWarning, catch_issues
+from cognite.neat._issues import IssueList, MultiValueError, NeatError, NeatWarning, catch_issues
 from cognite.neat._issues.errors import NeatTypeError
 from cognite.neat._rules._shared import (
     InputRules,
@@ -18,6 +18,8 @@ from cognite.neat._rules.models import (
     InformationInputRules,
     InformationRules,
 )
+from cognite.neat._rules.models.dms import DMSValidation
+from cognite.neat._rules.models.information import InformationValidation
 
 from ._base import RulesTransformer
 
@@ -26,9 +28,11 @@ class VerificationTransformer(RulesTransformer[T_InputRules, T_VerifiedRules], A
     """Base class for all verification transformers."""
 
     _rules_cls: type[T_VerifiedRules]
+    _validation_cls: type
 
-    def __init__(self, errors: Literal["raise", "continue"]) -> None:
+    def __init__(self, errors: Literal["raise", "continue"], validate: bool = True) -> None:
         self.errors = errors
+        self.validate = validate
 
     def transform(self, rules: T_InputRules | OutRules[T_InputRules]) -> MaybeRules[T_VerifiedRules]:
         issues = IssueList()
@@ -39,7 +43,18 @@ class VerificationTransformer(RulesTransformer[T_InputRules, T_VerifiedRules], A
         verified_rules: T_VerifiedRules | None = None
         with catch_issues(issues, NeatError, NeatWarning, error_args) as future:
             rules_cls = self._get_rules_cls(in_)
-            verified_rules = rules_cls.model_validate(in_.dump())  # type: ignore[assignment]
+            dumped = in_.dump()
+            verified_rules = rules_cls.model_validate(dumped)  # type: ignore[assignment]
+            if self.validate:
+                validation_cls = self._get_validation_cls(verified_rules)  # type: ignore[arg-type]
+                validation_issues = validation_cls(verified_rules).validate()
+                # We need to trigger warnings are raise exceptions such that they are caught by the context manager
+                # and processed with the read context
+                if validation_issues.warnings:
+                    validation_issues.trigger_warnings()
+                if validation_issues.has_errors:
+                    verified_rules = None
+                    raise MultiValueError(validation_issues.errors)
 
         if (future.result == "failure" or issues.has_errors or verified_rules is None) and self.errors == "raise":
             raise issues.as_errors()
@@ -51,17 +66,22 @@ class VerificationTransformer(RulesTransformer[T_InputRules, T_VerifiedRules], A
     def _get_rules_cls(self, in_: T_InputRules) -> type[T_VerifiedRules]:
         return self._rules_cls
 
+    def _get_validation_cls(self, rules: T_VerifiedRules) -> type:
+        return self._validation_cls
+
 
 class VerifyDMSRules(VerificationTransformer[DMSInputRules, DMSRules]):
     """Class to verify DMS rules."""
 
     _rules_cls = DMSRules
+    _validation_cls = DMSValidation
 
 
 class VerifyInformationRules(VerificationTransformer[InformationInputRules, InformationRules]):
     """Class to verify Information rules."""
 
     _rules_cls = InformationRules
+    _validation_cls = InformationValidation
 
 
 class VerifyAnyRules(VerificationTransformer[InputRules, VerifiedRules]):
@@ -74,3 +94,11 @@ class VerifyAnyRules(VerificationTransformer[InputRules, VerifiedRules]):
             return DMSRules
         else:
             raise NeatTypeError(f"Unsupported rules type: {type(in_)}")
+
+    def _get_validation_cls(self, rules: T_VerifiedRules) -> type:
+        if isinstance(rules, InformationRules):
+            return InformationValidation
+        elif isinstance(rules, DMSRules):
+            return DMSValidation
+        else:
+            raise NeatTypeError(f"Unsupported rules type: {type(rules)}")

cognite/neat/_session/_base.py

@@ -5,11 +5,14 @@ from cognite.client import CogniteClient
 from cognite.client import data_modeling as dm
 
 from cognite.neat import _version
+from cognite.neat._client import NeatClient
 from cognite.neat._issues import IssueList, catch_issues
 from cognite.neat._issues.errors import RegexViolationError
 from cognite.neat._rules import importers
 from cognite.neat._rules._shared import ReadRules, VerifiedRules
 from cognite.neat._rules.models import DMSRules
+from cognite.neat._rules.models.dms import DMSValidation
+from cognite.neat._rules.models.information import InformationValidation
 from cognite.neat._rules.models.information._rules import InformationRules
 from cognite.neat._rules.models.information._rules_input import InformationInputRules
 from cognite.neat._rules.transformers import ConvertToRules, VerifyAnyRules
@@ -18,10 +21,11 @@ from cognite.neat._store._provenance import (
     INSTANCES_ENTITY,
     Change,
 )
-from cognite.neat._utils.auth import _CLIENT_NAME
 
 from ._collector import _COLLECTOR, Collector
+from ._drop import DropAPI
 from ._inspect import InspectAPI
+from ._mapping import MappingAPI
 from ._prepare import PrepareAPI
 from ._read import ReadAPI
 from ._set import SetAPI
@@ -41,19 +45,19 @@ class NeatSession:
         verbose: bool = True,
         load_engine: Literal["newest", "cache", "skip"] = "cache",
     ) -> None:
-        self._client = client
+        self._client = NeatClient(client) if client else None
         self._verbose = verbose
         self._state = SessionState(store_type=storage)
-        self.read = ReadAPI(self._state, client, verbose)
-        self.to = ToAPI(self._state, client, verbose)
-        self.prepare = PrepareAPI(self._state, verbose)
+        self.read = ReadAPI(self._state, self._client, verbose)
+        self.to = ToAPI(self._state, self._client, verbose)
+        self.prepare = PrepareAPI(self._client, self._state, verbose)
         self.show = ShowAPI(self._state)
         self.set = SetAPI(self._state, verbose)
         self.inspect = InspectAPI(self._state)
+        self.mapping = MappingAPI(self._state)
+        self.drop = DropAPI(self._state)
         self.opt = OptAPI()
         self.opt._display()
-        if self._client is not None and self._client._config is not None:
-            self._client._config.client_name = _CLIENT_NAME
         if load_engine != "skip" and (engine_version := load_neat_engine(client, load_engine)):
             print(f"Neat Engine {engine_version} loaded.")
 
@@ -63,9 +67,21 @@ class NeatSession:
 
     def verify(self) -> IssueList:
         source_id, last_unverified_rule = self._state.data_model.last_unverified_rule
-        transformer = VerifyAnyRules("continue")
+        transformer = VerifyAnyRules("continue", validate=False)
         start = datetime.now(timezone.utc)
         output = transformer.try_transform(last_unverified_rule)
+        if isinstance(output.rules, DMSRules):
+            issues = DMSValidation(output.rules, self._client).validate()
+        elif isinstance(output.rules, InformationRules):
+            issues = InformationValidation(output.rules).validate()
+        else:
+            raise NeatSessionError("Unsupported rule type")
+        if issues.has_errors:
+            # This is up for discussion, but I think we should not return rules that
+            # only pass the verification but not the validation.
+            output.rules = None
+        output.issues.extend(issues)
+
         end = datetime.now(timezone.utc)
 
         if output.rules:

cognite/neat/_session/_drop.py

@@ -0,0 +1,35 @@
+from rdflib import URIRef
+
+from ._state import SessionState
+from .exceptions import session_class_wrapper
+
+try:
+    from rich import print
+except ImportError:
+    ...
+
+
+@session_class_wrapper
+class DropAPI:
+    def __init__(self, state: SessionState):
+        self._state = state
+
+    def instances(self, type: str | list[str]) -> None:
+        """Drop instances from the session.
+
+        Args:
+            type: The type of instances to drop.
+        """
+        type_list = type if isinstance(type, list) else [type]
+        uri_type_type = dict((v, k) for k, v in self._state.instances.store.queries.types.items())
+        selected_uri_by_type: dict[URIRef, str] = {}
+        for type_item in type_list:
+            if type_item not in uri_type_type:
+                print(f"Type {type_item} not found.")
+            selected_uri_by_type[uri_type_type[type_item]] = type_item
+
+        result = self._state.instances.store.queries.drop_types(list(selected_uri_by_type.keys()))
+
+        for type_uri, count in result.items():
+            print(f"Dropped {count} instances of type {selected_uri_by_type[type_uri]}")
+        return None

cognite/neat/_session/_inspect.py

@@ -11,6 +11,13 @@ from cognite.neat._utils.upload import UploadResult, UploadResultCore, UploadRes
 from ._state import SessionState
 from .exceptions import session_class_wrapper
 
+try:
+    from rich.markdown import Markdown as RichMarkdown
+
+    RICH_AVAILABLE = True
+except ImportError:
+    RICH_AVAILABLE = False
+
 
 @session_class_wrapper
 class InspectAPI:
@@ -61,14 +68,19 @@ class InspectIssues:
             closest_match = set(difflib.get_close_matches(search, unique_types))
             issues = IssueList([issue for issue in issues if type(issue).__name__ in closest_match])
 
+        issue_str = "\n".join(
+            [f"  * **{type(issue).__name__}**: {issue.as_message(include_type=False)}" for issue in issues]
+        )
+        markdown_str = f"### {len(issues)} issues found\n\n{issue_str}"
+
         if IN_NOTEBOOK:
             from IPython.display import Markdown, display
 
-            issue_str = "\n".join(
-                [f"  * **{type(issue).__name__}**: {issue.as_message(include_type=False)}" for issue in issues]
-            )
-            message = f"### {len(issues)} issues found\n\n{issue_str}"
-            display(Markdown(message))
+            display(Markdown(markdown_str))
+        elif RICH_AVAILABLE:
+            from rich import print
+
+            print(RichMarkdown(markdown_str))
 
         if return_dataframe:
             return issues.to_pandas()

cognite/neat/_session/_mapping.py

@@ -0,0 +1,39 @@
+from datetime import datetime, timezone
+
+from cognite.neat._rules.models.mapping import load_classic_to_core_mapping
+from cognite.neat._rules.transformers import RuleMapper
+from cognite.neat._store._provenance import Change
+
+from ._state import SessionState
+from .exceptions import session_class_wrapper
+
+
+@session_class_wrapper
+class MappingAPI:
+    def __init__(self, state: SessionState):
+        self._state = state
+
+    def classic_to_core(self, org_name: str) -> None:
+        """Map classic types to core types.
+
+        Note this automatically creates an extended CogniteCore model.
+
+        """
+        source_id, rules = self._state.data_model.last_verified_dms_rules
+
+        start = datetime.now(timezone.utc)
+        transformer = RuleMapper(load_classic_to_core_mapping(org_name, rules.metadata.space, rules.metadata.version))
+        output = transformer.transform(rules)
+        end = datetime.now(timezone.utc)
+
+        change = Change.from_rules_activity(
+            output,
+            transformer.agent,
+            start,
+            end,
+            "Mapping classic to core",
+            self._state.data_model.provenance.source_entity(source_id)
+            or self._state.data_model.provenance.target_entity(source_id),
+        )
+
+        self._state.data_model.write(output.rules, change)