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

cognite/neat/rules/models/data_types.py

@@ -1,30 +1,62 @@
+import re
 import sys
+import typing
 from datetime import date, datetime
 from typing import Any, ClassVar
 
 from cognite.client.data_classes import data_modeling as dms
-from pydantic import BaseModel, model_serializer, model_validator
+from cognite.client.data_classes.data_modeling.data_types import Enum as DMSEnum
+from pydantic import BaseModel, Field, model_serializer, model_validator
 from pydantic.functional_validators import ModelWrapValidatorHandler
 
+from cognite.neat.rules.models.entities._single_value import ClassEntity, UnitEntity
+from cognite.neat.utils.regex_patterns import SPLIT_ON_COMMA_PATTERN, SPLIT_ON_EQUAL_PATTERN
+
 if sys.version_info >= (3, 11):
     from typing import Self
 else:
     from typing_extensions import Self
 
+# This patterns matches a string that is a data type, with optional content in parentheses.
+# For example, it matches "float(unit=power:megaw)" as name="float" and content="unit=power:megaw"
+_DATATYPE_PATTERN = re.compile(r"^(?P<name>[^(:]*)(\((?P<content>.+)\))?$")
 
-class DataType(BaseModel):
-    # These are necessary for Pydantic to work
-    # pydantic gets confused as we have no fields.
-    __pydantic_extra__ = None
-    __pydantic_fields_set__ = set()
-    __pydantic_private__ = {}
 
-    name: ClassVar[str]
+class DataType(BaseModel):
     python: ClassVar[type]
     dms: ClassVar[type[dms.PropertyType]]
     graphql: ClassVar[str]
     xsd: ClassVar[str]
     sql: ClassVar[str]
+    _dms_loaded: bool = False
+    # Repeat all here, just to make mypy happy
+    name: typing.Literal[
+        "boolean",
+        "token",
+        "float",
+        "double",
+        "integer",
+        "nonPositiveInteger",
+        "nonNegativeInteger",
+        "long",
+        "negativeInteger",
+        "string",
+        "langString",
+        "anyURI",
+        "normalizedString",
+        "dateTime",
+        "timestamp",
+        "dateTimeStamp",
+        "date",
+        "plainLiteral",
+        "Literal",
+        "timeseries",
+        "file",
+        "sequence",
+        "json",
+        "enum",
+        "",
+    ] = ""
 
     @classmethod
     def load(cls, data: Any) -> Self:
@@ -35,15 +67,31 @@ class DataType(BaseModel):
 
     @model_validator(mode="wrap")
     def _load(cls, value: Any, handler: ModelWrapValidatorHandler["DataType"]) -> Any:
-        if isinstance(value, cls | dict):
-            return value
-        elif isinstance(value, str):
-            value_standardized = value.casefold()
-            if cls_ := _DATA_TYPE_BY_DMS_TYPE.get(value_standardized):
-                return cls_()
-            elif cls_ := _DATA_TYPE_BY_NAME.get(value_standardized):
-                return cls_()
-            raise ValueError(f"Unknown literal type: {value}") from None
+        if cls is not DataType or isinstance(value, DataType):
+            # This is a subclass, let the subclass handle it
+            return handler(value)
+        elif isinstance(value, str) and (match := _DATATYPE_PATTERN.match(value)):
+            name = match.group("name").casefold()
+            cls_: type[DataType]
+            if name in _DATA_TYPE_BY_DMS_TYPE:
+                cls_ = _DATA_TYPE_BY_DMS_TYPE[name]
+                dms_loaded = True
+            elif name in _DATA_TYPE_BY_NAME:
+                cls_ = _DATA_TYPE_BY_NAME[name]
+                dms_loaded = False
+            else:
+                raise ValueError(f"Unknown data type: {value}") from None
+            extra_args: dict[str, Any] = {}
+            if content := match.group("content"):
+                extra_args = dict(
+                    SPLIT_ON_EQUAL_PATTERN.split(pair.strip()) for pair in SPLIT_ON_COMMA_PATTERN.split(content)
+                )
+                # Todo? Raise warning if extra_args contains keys that are not in the model fields
+            instance = cls_(**extra_args)
+            created = handler(instance)
+            # Private attributes are not validated or set. We need to set it manually
+            created._dms_loaded = dms_loaded
+            return created
         raise ValueError(f"Cannot load {cls.__name__} from {value}")
 
     @model_serializer(when_used="unless-none", return_type=str)
@@ -51,227 +99,285 @@ class DataType(BaseModel):
         return str(self)
 
     def __str__(self) -> str:
-        return self.name
+        if self._dms_loaded:
+            base = self.dms._type
+        else:
+            base = self.model_fields["name"].default
+        return self._suffix_extra_args(base)
+
+    def _suffix_extra_args(self, base: str) -> str:
+        extra_fields: dict[str, str] = {}
+        for field_id, field_ in self.model_fields.items():
+            if field_id == "name":
+                continue
+            value = getattr(self, field_id)
+            if value is None:
+                continue
+            extra_fields[field_.alias or field_id] = str(value)
+        if extra_fields:
+            content = ",".join(f"{key}={value}" for key, value in sorted(extra_fields.items(), key=lambda x: x[0]))
+            return f"{base}({content})"
+        return base
 
     def __eq__(self, other: Any) -> bool:
-        return isinstance(other, type(self))
+        return isinstance(other, type(self)) and str(self) == str(other)
 
     def __hash__(self) -> int:
         return hash(str(self))
 
     @classmethod
     def is_data_type(cls, value: str) -> bool:
-        return value.casefold() in _DATA_TYPE_BY_NAME or value.casefold() in _DATA_TYPE_BY_DMS_TYPE
+        if match := _DATATYPE_PATTERN.match(value):
+            name = match.group("name").casefold()
+            return name in _DATA_TYPE_BY_NAME or name in _DATA_TYPE_BY_DMS_TYPE
+        return False
 
 
 class Boolean(DataType):
-    name = "boolean"
     python = bool
     dms = dms.Boolean
     graphql = "Boolean"
     xsd = "boolean"
     sql = "BOOLEAN"
+    name: typing.Literal["boolean"] = "boolean"
 
 
 class Float(DataType):
-    name = "float"
     python = float
     dms = dms.Float32
     graphql = "Float"
     xsd = "float"
     sql = "FLOAT"
 
+    name: typing.Literal["float"] = "float"
+    unit: UnitEntity | None = None
+
 
 class Double(DataType):
-    name = "double"
     python = float
     dms = dms.Float64
     graphql = "Float"
     xsd = "double"
     sql = "FLOAT"
 
+    name: typing.Literal["double"] = "double"
+    unit: UnitEntity | None = None
+
 
 class Integer(DataType):
-    name = "integer"
     python = int
     dms = dms.Int32
     graphql = "Int"
     xsd = "integer"
     sql = "INTEGER"
 
+    name: typing.Literal["integer"] = "integer"
+
 
 class NonPositiveInteger(DataType):
-    name = "nonPositiveInteger"
     python = int
     dms = dms.Int32
     graphql = "Int"
     xsd = "nonPositiveInteger"
     sql = "INTEGER"
 
+    name: typing.Literal["nonPositiveInteger"] = "nonPositiveInteger"
+
 
 class NonNegativeInteger(DataType):
-    name = "nonNegativeInteger"
     python = int
     dms = dms.Int32
     graphql = "Int"
     xsd = "nonNegativeInteger"
     sql = "INTEGER"
 
+    name: typing.Literal["nonNegativeInteger"] = "nonNegativeInteger"
+
 
 class NegativeInteger(DataType):
-    name = "negativeInteger"
     python = int
     dms = dms.Int32
     graphql = "Int"
     xsd = "negativeInteger"
     sql = "INTEGER"
 
+    name: typing.Literal["negativeInteger"] = "negativeInteger"
+
 
 class Long(DataType):
-    name = "long"
     python = int
     dms = dms.Int64
     graphql = "Int"
     xsd = "long"
     sql = "BIGINT"
 
+    name: typing.Literal["long"] = "long"
+
 
 class String(DataType):
-    name = "string"
     python = str
     dms = dms.Text
     graphql = "String"
     xsd = "string"
     sql = "STRING"
 
+    name: typing.Literal["string"] = "string"
+
 
 class LangString(DataType):
-    name = "langString"
     python = str
     dms = dms.Text
     graphql = "String"
     xsd = "string"
     sql = "STRING"
 
+    name: typing.Literal["langString"] = "langString"
+
 
 class AnyURI(DataType):
-    name = "anyURI"
     python = str
     dms = dms.Text
     graphql = "String"
     xsd = "anyURI"
     sql = "STRING"
 
+    name: typing.Literal["anyURI"] = "anyURI"
+
 
 class NormalizedString(DataType):
-    name = "normalizedString"
     python = str
     dms = dms.Text
     graphql = "String"
     xsd = "normalizedString"
     sql = "STRING"
 
+    name: typing.Literal["normalizedString"] = "normalizedString"
+
 
 class Token(DataType):
-    name = "token"
     python = str
     dms = dms.Text
     graphql = "String"
     xsd = "string"
     sql = "STRING"
 
+    name: typing.Literal["token"] = "token"
+
 
 class DateTime(DataType):
-    name = "dateTime"
     python = datetime
     dms = dms.Timestamp
     graphql = "Timestamp"
     xsd = "dateTimeStamp"
     sql = "TIMESTAMP"
 
+    name: typing.Literal["dateTime"] = "dateTime"
+
 
 class Timestamp(DataType):
-    name = "timestamp"
     python = datetime
     dms = dms.Timestamp
     graphql = "Timestamp"
     xsd = "dateTimeStamp"
     sql = "TIMESTAMP"
 
+    name: typing.Literal["timestamp"] = "timestamp"
+
 
 class DateTimeStamp(DataType):
-    name = "dateTimeStamp"
     python = datetime
     dms = dms.Timestamp
     graphql = "Timestamp"
     xsd = "dateTimeStamp"
     sql = "TIMESTAMP"
 
+    name: typing.Literal["dateTimeStamp"] = "dateTimeStamp"
+
 
 class Date(DataType):
-    name = "date"
     python = date
     dms = dms.Date
     graphql = "String"
     xsd = "date"
     sql = "DATE"
 
+    name: typing.Literal["date"] = "date"
+
 
 class PlainLiteral(DataType):
-    name = "PlainLiteral"
     python = str
     dms = dms.Text
     graphql = "String"
     xsd = "plainLiteral"
     sql = "STRING"
 
+    name: typing.Literal["plainLiteral"] = "plainLiteral"
+
 
 class Literal(DataType):
-    name = "Literal"
     python = str
     dms = dms.Text
     graphql = "String"
     xsd = "string"
     sql = "STRING"
 
+    name: typing.Literal["Literal"] = "Literal"
+
 
 class Timeseries(DataType):
-    name = "timeseries"
     python = str
     dms = dms.TimeSeriesReference
     graphql = "TimeSeries"
     xsd = "string"
     sql = "STRING"
 
+    name: typing.Literal["timeseries"] = "timeseries"
+
 
 class File(DataType):
-    name = "file"
     python = str
     dms = dms.FileReference
     graphql = "File"
     xsd = "string"
     sql = "STRING"
 
+    name: typing.Literal["file"] = "file"
+
 
 class Sequence(DataType):
-    name = "sequence"
     python = str
     dms = dms.SequenceReference
     graphql = "Sequence"
     xsd = "string"
     sql = "STRING"
 
+    name: typing.Literal["sequence"] = "sequence"
+
 
 class Json(DataType):
-    name = "json"
     python = dict
     dms = dms.Json
     graphql = "Json"
     xsd = "json"
     sql = "STRING"
 
+    name: typing.Literal["json"] = "json"
+
+
+class Enum(DataType):
+    python = str
+    dms = DMSEnum
+    graphql = "Enum"
+    xsd = "string"
+    sql = "STRING"
+
+    name: typing.Literal["enum"] = "enum"
+
+    collection: ClassEntity
+    unknown_value: str | None = Field(None, alias="unknownValue")
+
 
-_DATA_TYPE_BY_NAME = {cls.name.casefold(): cls for cls in DataType.__subclasses__()}
+_DATA_TYPE_BY_NAME = {cls.model_fields["name"].default.casefold(): cls for cls in DataType.__subclasses__()}
 _seen = set()
 _DATA_TYPE_BY_DMS_TYPE = {
     cls.dms._type.casefold(): cls

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

@@ -1,5 +1,13 @@
-from ._rules import DMSContainer, DMSMetadata, DMSProperty, DMSRules, DMSView
-from ._rules_input import DMSContainerInput, DMSMetadataInput, DMSPropertyInput, DMSRulesInput, DMSViewInput
+from ._rules import DMSContainer, DMSEnum, DMSMetadata, DMSNode, DMSProperty, DMSRules, DMSView
+from ._rules_input import (
+    DMSInputContainer,
+    DMSInputEnum,
+    DMSInputMetadata,
+    DMSInputNode,
+    DMSInputProperty,
+    DMSInputRules,
+    DMSInputView,
+)
 from ._schema import DMSSchema, PipelineSchema
 
 __all__ = [
@@ -9,10 +17,14 @@ __all__ = [
     "DMSView",
     "DMSProperty",
     "DMSContainer",
+    "DMSNode",
+    "DMSEnum",
     "PipelineSchema",
-    "DMSRulesInput",
-    "DMSMetadataInput",
-    "DMSViewInput",
-    "DMSPropertyInput",
-    "DMSContainerInput",
+    "DMSInputRules",
+    "DMSInputMetadata",
+    "DMSInputView",
+    "DMSInputProperty",
+    "DMSInputContainer",
+    "DMSInputNode",
+    "DMSInputEnum",
 ]

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,11 +318,27 @@ 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
@@ -463,7 +500,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 +509,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 +523,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.
@@ -498,6 +536,7 @@ class _DMSExporter:
                     f"Debug Info, Invalid valueType reverse connection: {prop.model_dump_json()}"
                 )
             reverse_prop: DMSProperty | None = None
+            edge_source = None
             if reverse_prop_id is not None:
                 reverse_prop = next(
                     (
@@ -507,6 +546,12 @@ class _DMSExporter:
                     ),
                     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 +565,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,6 +575,7 @@ 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":
                 reverse_direct_cls = (