cognite-neat 0.88.2__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",
 ]