[NEAT-183] Apply new entities (#420)

* refactor: first pass

* refactor: better implementation

* refactor; handle dict

* refactor; handle dict

* refactor; serializable entity

* refactor; handle default space and version

* refactor: backwards compatible

* refactor: backwards compatible

* refactor: fix issue with id

* feat: ParentEntityList

* entities: backwards compatability

* replace: ClassEntity and ParentClassEntity

* refactor: replace information rules entities and data types

* refactor; domain rules

* refator: first pass on full replace

* refactor; working towards new data types and enteties

* refactor: Isolate problem

* refactor; replace rdflib.HTTP

* refactor; reload alias

* refactor: addd dms write rules

* tests: upgraded test to new input class

* refactor: added load methods

* tests: updated dms tests -1

* refactor; testing of union parse order

* refactor: update reference to correct parsing

* tests: All information rules passing

* tests: Pass all dms architect

* fix: handle missing source

* refactor: fix dms2rules

* refactor: correct dumping

* fix: export URL correctly

* fix: yaml importer

* tests: exporters passing

* refactor: pass last rules test

* Linting and static code checks

* refactor: update bad imports

* refactor: update bad imports

* refactor; happier mypy

* refactor: happier mypy

* refactor: happier mypy

* style: happy mypy

* style: pre-commit

* refactor; fix failed fix

* refactor: cleanup round1

* tests: moved test over

* refactor: cleanup round2

* refactor: cleanup round 3

* refactor: removed unused regex

* refactor: require prefix reference

* refactor: replace xsd:

* refactor: cleanup

---------

Co-authored-by: doctrino <doctrino@users.noreply.github.com>

cognite/neat/graph/extractors/_mock_graph_generator.py

@@ -14,9 +14,10 @@
 
 from cognite.neat.graph.models import Triple
 from cognite.neat.rules.analysis import InformationArchitectRulesAnalysis
+from cognite.neat.rules.models.data_types import DataType
+from cognite.neat.rules.models.entities import ClassEntity, EntityTypes
 from cognite.neat.rules.models.rules import DMSRules, InformationRules
 from cognite.neat.rules.models.rules._information_rules import InformationProperty
-from cognite.neat.rules.models.rules._types import ClassEntity, EntityTypes, XSDValueType
 from cognite.neat.utils.utils import remove_namespace
 
 from ._base import BaseExtractor
@@ -55,7 +56,7 @@ def __init__(
             }
         elif all(isinstance(key, str) for key in class_count.keys()):
             self.class_count = {
-                ClassEntity.from_raw(f"{self.rules.metadata.prefix}:{key}"): value for key, value in class_count.items()
+                ClassEntity.load(f"{self.rules.metadata.prefix}:{key}"): value for key, value in class_count.items()
             }
         elif all(isinstance(key, ClassEntity) for key in class_count.keys()):
             self.class_count = cast(dict[ClassEntity, int], class_count)
@@ -147,7 +148,7 @@ def generate_triples(
     triples: list[Triple] = []
     for class_ in class_count:
         triples += [
-            (class_instance_id, RDF.type, URIRef(namespace[class_.suffix]))
+            (class_instance_id, RDF.type, URIRef(namespace[str(class_.suffix)]))
             for class_instance_id in instance_ids[class_]
         ]
 
@@ -246,7 +247,7 @@ def _remove_non_requested_sym_pairs(class_linkage: pd.DataFrame, class_count: di
 
 
 def _generate_mock_data_property_triples(
-    instance_ids: list[URIRef], property_: str, namespace: Namespace, value_type: XSDValueType
+    instance_ids: list[URIRef], property_: str, namespace: Namespace, value_type: DataType
 ) -> list[tuple[URIRef, URIRef, Literal]]:
     """Generates triples for data properties."""
 
@@ -338,7 +339,7 @@ def _generate_triples_per_class(
                 instance_ids[class_],
                 property_.property_,
                 namespace,
-                cast(XSDValueType, property_.value_type),
+                cast(DataType, property_.value_type),
             )
 
         elif property_.type_ == EntityTypes.object_property:

cognite/neat/rules/analysis/_base.py

@@ -3,7 +3,7 @@
 from typing import Generic, TypeVar
 
 from cognite.neat.rules._shared import Rules
-from cognite.neat.rules.models.rules._types import ClassEntity
+from cognite.neat.rules.models.entities import ClassEntity
 
 if sys.version_info >= (3, 11):
     from enum import StrEnum

cognite/neat/rules/analysis/_information_rules.py

@@ -7,10 +7,10 @@
 import pandas as pd
 from pydantic import ValidationError
 
+from cognite.neat.rules.models.entities import ClassEntity, EntityTypes, ParentClassEntity, ReferenceEntity
 from cognite.neat.rules.models.rdfpath import TransformationRuleType
 from cognite.neat.rules.models.rules._base import SchemaCompleteness
 from cognite.neat.rules.models.rules._information_rules import InformationClass, InformationProperty, InformationRules
-from cognite.neat.rules.models.rules._types._base import ClassEntity, EntityTypes, ParentClassEntity, ReferenceEntity
 from cognite.neat.utils.utils import get_inheritance_path
 
 from ._base import BaseAnalysis, DataModelingScenario
@@ -126,7 +126,7 @@ def classes_with_properties(
         return class_property_pairs
 
     def class_inheritance_path(self, class_: ClassEntity | str, use_reference: bool = False) -> list[ClassEntity]:
-        class_ = class_ if isinstance(class_, ClassEntity) else ClassEntity.from_raw(class_)
+        class_ = class_ if isinstance(class_, ClassEntity) else ClassEntity.load(class_)
         class_parent_pairs = self.class_parent_pairs(use_reference)
         return get_inheritance_path(class_, class_parent_pairs)
 
@@ -328,7 +328,7 @@ def as_class_dict(self) -> dict[str, InformationClass]:
         """This is to simplify access to classes through dict."""
         class_dict: dict[str, InformationClass] = {}
         for definition in self.rules.classes:
-            class_dict[definition.class_.suffix] = definition
+            class_dict[str(definition.class_.suffix)] = definition
         return class_dict
 
     def subset_rules(self, desired_classes: set[ClassEntity], use_reference: bool = False) -> InformationRules:
@@ -398,7 +398,7 @@ def subset_rules(self, desired_classes: set[ClassEntity], use_reference: bool =
 
         logging.info(f"Reducing data model to only include the following classes: {possible_classes}")
         for class_ in possible_classes:
-            reduced_data_model["classes"].append(class_as_dict[class_.suffix])
+            reduced_data_model["classes"].append(class_as_dict[str(class_.suffix)])
 
         class_property_pairs = self.classes_with_properties(consider_inheritance=False)
 

cognite/neat/rules/exporters/_rules2ontology.py

@@ -11,14 +11,15 @@
 from cognite.neat.constants import DEFAULT_NAMESPACE as NEAT_NAMESPACE
 from cognite.neat.rules import exceptions
 from cognite.neat.rules.analysis import InformationArchitectRulesAnalysis
+from cognite.neat.rules.models.data_types import DataType
+from cognite.neat.rules.models.entities import ClassEntity, EntityTypes
 from cognite.neat.rules.models.rules import DMSRules
 from cognite.neat.rules.models.rules._information_rules import (
     InformationClass,
     InformationMetadata,
     InformationProperty,
     InformationRules,
 )
-from cognite.neat.rules.models.rules._types import XSD_VALUE_TYPE_MAPPINGS, EntityTypes
 from cognite.neat.utils.utils import generate_exception_report, remove_namespace
 
 from ._base import BaseExporter
@@ -114,7 +115,7 @@ def from_rules(cls, input_rules: Rules) -> Self:
             ],
             shapes=[
                 SHACLNodeShape.from_rules(
-                    class_dict[class_.suffix],
+                    class_dict[str(class_.suffix)],
                     list(properties.values()),
                     rules.metadata.namespace,
                 )
@@ -249,15 +250,15 @@ def from_class(cls, definition: InformationClass, namespace: Namespace, prefixes
             sub_class_of = []
             for parent_class in definition.parent:
                 try:
-                    sub_class_of.append(prefixes[parent_class.prefix][parent_class.suffix])
+                    sub_class_of.append(prefixes[str(parent_class.prefix)][str(parent_class.suffix)])
                 except KeyError:
-                    sub_class_of.append(namespace[parent_class.suffix])
+                    sub_class_of.append(namespace[str(parent_class.suffix)])
         else:
             sub_class_of = None
 
         return cls(
-            id_=namespace[definition.class_.suffix],
-            label=definition.name or definition.class_.suffix,
+            id_=namespace[str(definition.class_.suffix)],
+            label=definition.name or str(definition.class_.suffix),
             comment=definition.description,
             sub_class_of=sub_class_of,
             namespace=namespace,
@@ -324,12 +325,14 @@ def from_list_of_properties(cls, definitions: list[InformationProperty], namespa
         )
         for definition in definitions:
             owl_property.type_.add(OWL[definition.type_])
-            owl_property.range_.add(
-                XSD[definition.value_type.suffix]
-                if definition.value_type.suffix in XSD_VALUE_TYPE_MAPPINGS
-                else namespace[definition.value_type.suffix]
-            )
-            owl_property.domain.add(namespace[definition.class_.suffix])
+
+            if isinstance(definition.value_type, DataType):
+                owl_property.range_.add(XSD[definition.value_type.xsd])
+            elif isinstance(definition.value_type, ClassEntity):
+                owl_property.range_.add(namespace[str(definition.value_type.suffix)])
+            else:
+                raise ValueError(f"Value type {definition.value_type} is not supported")
+            owl_property.domain.add(namespace[str(definition.class_.suffix)])
             owl_property.label.add(definition.name or definition.property_)
             if definition.description:
                 owl_property.comment.add(definition.description)
@@ -493,12 +496,12 @@ def from_rules(
         cls, class_definition: InformationClass, property_definitions: list[InformationProperty], namespace: Namespace
     ) -> "SHACLNodeShape":
         if class_definition.parent:
-            parent = [namespace[parent.suffix + "Shape"] for parent in class_definition.parent]
+            parent = [namespace[str(parent.suffix) + "Shape"] for parent in class_definition.parent]
         else:
             parent = None
         return cls(
-            id_=namespace[f"{class_definition.class_.suffix}Shape"],
-            target_class=namespace[class_definition.class_.suffix],
+            id_=namespace[f"{class_definition.class_.suffix!s}Shape"],
+            target_class=namespace[str(class_definition.class_.suffix)],
             parent=parent,
             property_shapes=[SHACLPropertyShape.from_property(prop, namespace) for prop in property_definitions],
             namespace=namespace,
@@ -552,8 +555,8 @@ def from_property(cls, definition: InformationProperty, namespace: Namespace) ->
             node_kind=SHACL.IRI if definition.type_ == EntityTypes.object_property else SHACL.Literal,
             expected_value_type=(
                 namespace[f"{definition.value_type.suffix}Shape"]
-                if definition.type_ == EntityTypes.object_property
-                else XSD[definition.value_type.suffix]
+                if isinstance(definition.value_type, ClassEntity)
+                else XSD[definition.value_type.xsd]
             ),
             min_count=definition.min_count,
             max_count=(

cognite/neat/rules/exporters/_validation.py

@@ -26,7 +26,7 @@ def are_entity_names_dms_compliant(
     flag: bool = True
     with warnings.catch_warnings(record=True) as validation_warnings:
         for class_ in rules.classes:
-            if not re.match(VIEW_ID_COMPLIANCE_REGEX, class_.class_.suffix):
+            if not re.match(VIEW_ID_COMPLIANCE_REGEX, str(class_.class_.suffix)):
                 warnings.warn(
                     exceptions.EntityIDNotDMSCompliant(
                         "Class", class_.class_.versioned_id, f"[Classes/Class/{class_.class_.versioned_id}]"
@@ -38,7 +38,7 @@ def are_entity_names_dms_compliant(
 
         for row, property_ in enumerate(rules.properties):
             # check class id which would resolve as view/container id
-            if not re.match(VIEW_ID_COMPLIANCE_REGEX, property_.class_.suffix):
+            if not re.match(VIEW_ID_COMPLIANCE_REGEX, str(property_.class_.suffix)):
                 warnings.warn(
                     exceptions.EntityIDNotDMSCompliant(
                         "Class", property_.class_.versioned_id, f"[Properties/Class/{row}]"

cognite/neat/rules/importers/_dms2rules.py

@@ -10,6 +10,14 @@
 from cognite.neat.rules import issues
 from cognite.neat.rules.importers._base import BaseImporter, Rules
 from cognite.neat.rules.issues import IssueList
+from cognite.neat.rules.models.data_types import DataType
+from cognite.neat.rules.models.entities import (
+    ClassEntity,
+    ContainerEntity,
+    DMSUnknownEntity,
+    ViewEntity,
+    ViewPropertyEntity,
+)
 from cognite.neat.rules.models.rules import DMSRules, DMSSchema, RoleTypes
 from cognite.neat.rules.models.rules._base import ExtensionCategory, SchemaCompleteness
 from cognite.neat.rules.models.rules._dms_architect_rules import (
@@ -19,15 +27,6 @@
     DMSView,
     SheetList,
 )
-from cognite.neat.rules.models.rules._types import (
-    ClassEntity,
-    ContainerEntity,
-    DMSValueType,
-    Undefined,
-    Unknown,
-    ViewEntity,
-    ViewPropEntity,
-)
 
 
 class DMSImporter(BaseImporter):
@@ -136,21 +135,21 @@ def to_rules(
                             raise NotImplementedError(f"Constraint type {type(constraint_obj)} not implemented")
 
                     if isinstance(container_prop.type, dm.DirectRelation):
-                        direct_value_type: str | ViewEntity | DMSValueType
+                        direct_value_type: str | ViewEntity | DataType | DMSUnknownEntity
                         if prop.source is None:
                             issue_list.append(
                                 issues.importing.UnknownValueTypeWarning(class_entity.versioned_id, prop_id)
                             )
-                            direct_value_type = ViewPropEntity(prefix=Undefined, suffix=Unknown)
+                            direct_value_type = DMSUnknownEntity()
                         else:
-                            direct_value_type = ViewPropEntity.from_id(prop.source)
+                            direct_value_type = ViewEntity.from_id(prop.source)
 
                         dms_property = DMSProperty(
                             class_=class_entity,
                             property_=prop_id,
                             description=prop.description,
                             name=prop.name,
-                            value_type=cast(ViewPropEntity | DMSValueType, direct_value_type),
+                            value_type=direct_value_type,
                             relation="direct",
                             nullable=container_prop.nullable,
                             default=container_prop.default_value,
@@ -168,7 +167,7 @@ def to_rules(
                             property_=prop_id,
                             description=prop.description,
                             name=prop.name,
-                            value_type=cast(ViewPropEntity | DMSValueType, container_prop.type._type),
+                            value_type=cast(ViewPropertyEntity | DataType, container_prop.type._type),
                             nullable=container_prop.nullable,
                             is_list=container_prop.type.is_list,
                             default=container_prop.default_value,
@@ -180,7 +179,7 @@ def to_rules(
                             constraint=unique_constraints or None,
                         )
                 elif isinstance(prop, dm.MultiEdgeConnectionApply):
-                    view_entity = ViewPropEntity.from_id(prop.source)
+                    view_entity = ViewEntity.from_id(prop.source)
                     dms_property = DMSProperty(
                         class_=ClassEntity(prefix=view.space, suffix=view.external_id, version=view.version),
                         property_=prop_id,

cognite/neat/rules/importers/_dtdl2rules/dtdl_converter.py

@@ -20,13 +20,9 @@
     TelemetryV2,
 )
 from cognite.neat.rules.issues import IssueList, ValidationIssue
+from cognite.neat.rules.models.data_types import _DATA_TYPE_BY_NAME, DataType, Json, String
+from cognite.neat.rules.models.entities import ClassEntity, ParentClassEntity
 from cognite.neat.rules.models.rules._information_rules import InformationClass, InformationProperty
-from cognite.neat.rules.models.rules._types import (
-    XSD_VALUE_TYPE_MAPPINGS,
-    ClassEntity,
-    ParentClassEntity,
-    XSDValueType,
-)
 
 
 class _DTDLConverter:
@@ -92,7 +88,7 @@ def convert_interface(self, item: Interface, _: str | None) -> None:
             name=item.display_name,
             description=item.description,
             comment=item.comment,
-            parent=[ParentClassEntity.from_raw(parent.as_class_id()) for parent in item.extends or []] or None,
+            parent=[ParentClassEntity.load(parent.as_class_id()) for parent in item.extends or []] or None,
         )
         self.classes.append(class_)
         for sub_item_or_id in item.contents or []:
@@ -123,7 +119,7 @@ def convert_property(
             return None
 
         prop = InformationProperty(
-            class_=ClassEntity.from_raw(parent),
+            class_=ClassEntity.load(parent),
             property_=item.name,
             name=item.display_name,
             description=item.description,
@@ -175,7 +171,7 @@ def convert_command(self, item: Command | CommandV2, parent: str | None) -> None
         if value_type is None:
             return
         prop = InformationProperty(
-            class_=ClassEntity.from_raw(parent),
+            class_=ClassEntity.load(parent),
             property_=item.name,
             name=item.display_name,
             description=item.description,
@@ -195,7 +191,7 @@ def convert_component(self, item: Component, parent: str | None) -> None:
         if value_type is None:
             return
         prop = InformationProperty(
-            class_=ClassEntity.from_raw(parent),
+            class_=ClassEntity.load(parent),
             property_=item.name,
             name=item.display_name,
             description=item.description,
@@ -211,7 +207,7 @@ def convert_relationship(self, item: Relationship, parent: str | None) -> None:
             self._missing_parent_warning(item)
             return None
         if item.target is not None:
-            value_type: XSDValueType | ClassEntity
+            value_type: DataType | ClassEntity
             if item.target in self._item_by_id:
                 value_type = item.target.as_class_id()
             else:
@@ -223,10 +219,10 @@ def convert_relationship(self, item: Relationship, parent: str | None) -> None:
                         instance_id=item.target.model_dump(),
                     )
                 )
-                value_type = XSD_VALUE_TYPE_MAPPINGS["json"]
+                value_type = Json()
 
             prop = InformationProperty(
-                class_=ClassEntity.from_raw(parent),
+                class_=ClassEntity.load(parent),
                 property_=item.name,
                 name=item.display_name,
                 description=item.description,
@@ -275,13 +271,13 @@ def convert_object(self, item: Object, _: str | None) -> None:
 
     def schema_to_value_type(
         self, schema: Schema | Interface | DTMI | None, item: DTDLBase
-    ) -> XSDValueType | ClassEntity | None:
+    ) -> DataType | ClassEntity | None:
         input_type = self._item_by_id.get(schema) if isinstance(schema, DTMI) else schema
 
         if isinstance(input_type, Enum):
-            return XSD_VALUE_TYPE_MAPPINGS["string"]
-        elif isinstance(input_type, str) and input_type in XSD_VALUE_TYPE_MAPPINGS:
-            return XSD_VALUE_TYPE_MAPPINGS[input_type]
+            return String()
+        elif isinstance(input_type, str) and input_type.casefold() in _DATA_TYPE_BY_NAME:
+            return _DATA_TYPE_BY_NAME[input_type.casefold()]()
         elif isinstance(input_type, str):
             self.issues.append(
                 cognite.neat.rules.issues.importing.UnsupportedPropertyTypeError(
@@ -301,11 +297,11 @@ def schema_to_value_type(
                         instance_name=input_type.display_name,
                     )
                 )
-                return XSD_VALUE_TYPE_MAPPINGS["json"]
+                return Json()
             else:
                 if isinstance(input_type, Object):
                     self.convert_object(input_type, None)
-                return ClassEntity.from_raw(input_type.id_.as_class_id())
+                return ClassEntity.load(input_type.id_.as_class_id())
         else:
             self.issues.append(
                 issues.importing.UnknownPropertyWarning(

cognite/neat/rules/importers/_dtdl2rules/spec.py

@@ -14,7 +14,7 @@
 
 from pydantic import BaseModel, Field, field_validator, model_serializer, model_validator
 
-from cognite.neat.rules.models.rules._types import ClassEntity
+from cognite.neat.rules.models.entities import ClassEntity
 
 if TYPE_CHECKING:
     from pydantic.type_adapter import IncEx