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

cognite/neat/rules/models/asset/_rules_input.py

@@ -1,29 +1,33 @@
-from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any, cast, overload
+from typing import Any
 
-from cognite.neat.rules.models._base import _add_alias
+from rdflib import Namespace
+
+from cognite.neat.rules.models._base_input import InputComponent, InputRules
 from cognite.neat.rules.models.data_types import DataType
 from cognite.neat.rules.models.entities import (
     ClassEntity,
     MultiValueTypeInfo,
-    Unknown,
     UnknownEntity,
+    load_value_type,
 )
-from cognite.neat.rules.models.information._rules_input import InformationClassInput, InformationMetadataInput
+from cognite.neat.rules.models.information._rules_input import InformationInputClass, InformationInputMetadata
 
-from ._rules import AssetProperty, AssetRules
+from ._rules import AssetClass, AssetMetadata, AssetProperty, AssetRules
 
 
 @dataclass
-class AssetMetadataInput(InformationMetadataInput): ...
+class AssetInputMetadata(InformationInputMetadata):
+    @classmethod
+    def _get_verified_cls(cls) -> type[AssetMetadata]:
+        return AssetMetadata
 
 
 @dataclass
-class AssetPropertyInput:
-    class_: str
+class AssetInputProperty(InputComponent[AssetProperty]):
+    class_: ClassEntity | str
     property_: str
-    value_type: str
+    value_type: DataType | ClassEntity | MultiValueTypeInfo | UnknownEntity | str
     name: str | None = None
     description: str | None = None
     comment: str | None = None
@@ -34,133 +38,54 @@ class AssetPropertyInput:
     match_type: str | None = None
     transformation: str | None = None
     implementation: str | None = None
+    # Only used internally
+    inherited: bool = False
 
     @classmethod
-    @overload
-    def load(cls, data: None) -> None: ...
+    def _get_verified_cls(cls) -> type[AssetProperty]:
+        return AssetProperty
 
-    @classmethod
-    @overload
-    def load(cls, data: dict[str, Any]) -> "AssetPropertyInput": ...
+    def dump(self, default_prefix: str) -> dict[str, Any]:  # type: ignore[override]
+        output = super().dump()
+        output["Class"] = ClassEntity.load(self.class_, prefix=default_prefix)
+        output["Value Type"] = load_value_type(self.value_type, default_prefix)
+        return output
 
-    @classmethod
-    @overload
-    def load(cls, data: list[dict[str, Any]]) -> list["AssetPropertyInput"]: ...
 
+@dataclass
+class AssetInputClass(InformationInputClass):
     @classmethod
-    def load(
-        cls, data: dict[str, Any] | list[dict[str, Any]] | None
-    ) -> "AssetPropertyInput | list[AssetPropertyInput] | None":
-        if data is None:
-            return None
-        if isinstance(data, list) or (isinstance(data, dict) and isinstance(data.get("data"), list)):
-            items = cast(list[dict[str, Any]], data.get("data") if isinstance(data, dict) else data)
-            return [loaded for item in items if (loaded := cls.load(item)) is not None]
-
-        _add_alias(data, AssetProperty)
-        return cls(
-            class_=data.get("class_"),  # type: ignore[arg-type]
-            property_=data.get("property_"),  # type: ignore[arg-type]
-            name=data.get("name", None),
-            description=data.get("description", None),
-            comment=data.get("comment", None),
-            value_type=data.get("value_type"),  # type: ignore[arg-type]
-            min_count=data.get("min_count", None),
-            max_count=data.get("max_count", None),
-            default=data.get("default", None),
-            reference=data.get("reference", None),
-            match_type=data.get("match_type", None),
-            transformation=data.get("transformation", None),
-            implementation=data.get("implementation", None),
-        )
-
-    def dump(self, default_prefix: str) -> dict[str, Any]:
-        value_type: MultiValueTypeInfo | DataType | ClassEntity | UnknownEntity
-
-        # property holding xsd data type
-        # check if it is multi value type
-        if "|" in self.value_type:
-            value_type = MultiValueTypeInfo.load(self.value_type)
-            value_type.set_default_prefix(default_prefix)
-
-        elif DataType.is_data_type(self.value_type):
-            value_type = DataType.load(self.value_type)
-
-        # unknown value type
-        elif self.value_type == str(Unknown):
-            value_type = UnknownEntity()
-
-        # property holding link to class
-        else:
-            value_type = ClassEntity.load(self.value_type, prefix=default_prefix)
-
-        return {
-            "Class": ClassEntity.load(self.class_, prefix=default_prefix),
-            "Property": self.property_,
-            "Name": self.name,
-            "Description": self.description,
-            "Comment": self.comment,
-            "Value Type": value_type,
-            "Min Count": self.min_count,
-            "Max Count": self.max_count,
-            "Default": self.default,
-            "Reference": self.reference,
-            "Match Type": self.match_type,
-            "Transformation": self.transformation,
-            "Implementation": self.implementation,
-        }
-
-
-class AssetClassInput(InformationClassInput): ...
+    def _get_verified_cls(cls) -> type[AssetClass]:
+        return AssetClass
 
 
 @dataclass
-class AssetRulesInput:
-    metadata: AssetMetadataInput
-    properties: Sequence[AssetPropertyInput]
-    classes: Sequence[AssetClassInput]
-    last: "AssetRulesInput | AssetRules | None" = None
-    reference: "AssetRulesInput | AssetRules | None" = None
-
-    @classmethod
-    @overload
-    def load(cls, data: dict[str, Any]) -> "AssetRulesInput": ...
-
-    @classmethod
-    @overload
-    def load(cls, data: None) -> None: ...
+class AssetInputRules(InputRules[AssetRules]):
+    metadata: AssetInputMetadata
+    properties: list[AssetInputProperty]
+    classes: list[AssetInputClass]
+    prefixes: dict[str, Namespace] | None = None
+    last: "AssetInputRules | None" = None
+    reference: "AssetInputRules | None" = None
 
     @classmethod
-    def load(cls, data: dict | None) -> "AssetRulesInput | None":
-        if data is None:
-            return None
-        _add_alias(data, AssetRules)
-
-        return cls(
-            metadata=AssetMetadataInput.load(data.get("metadata")),  # type: ignore[arg-type]
-            properties=AssetPropertyInput.load(data.get("properties")),  # type: ignore[arg-type]
-            classes=InformationClassInput.load(data.get("classes")),  # type: ignore[arg-type]
-            last=AssetRulesInput.load(data.get("last")),
-            reference=AssetRulesInput.load(data.get("reference")),
-        )
-
-    def as_rules(self) -> AssetRules:
-        return AssetRules.model_validate(self.dump())
+    def _get_verified_cls(cls) -> type[AssetRules]:
+        return AssetRules
 
     def dump(self) -> dict[str, Any]:
         default_prefix = self.metadata.prefix
         reference: dict[str, Any] | None = None
-        if isinstance(self.reference, AssetRulesInput):
+        if isinstance(self.reference, AssetInputRules):
             reference = self.reference.dump()
         elif isinstance(self.reference, AssetRules):
             # We need to load through the AssetRulesInput to set the correct default space and version
-            reference = AssetRulesInput.load(self.reference.model_dump()).dump()
+            reference = AssetInputRules.load(self.reference.model_dump()).dump()
         last: dict[str, Any] | None = None
-        if isinstance(self.last, AssetRulesInput):
+        if isinstance(self.last, AssetInputRules):
             last = self.last.dump()
         elif isinstance(self.last, AssetRules):
             # We need to load through the AssetRulesInput to set the correct default space and version
-            last = AssetRulesInput.load(self.last.model_dump()).dump()
+            last = AssetInputRules.load(self.last.model_dump()).dump()
 
         return dict(
             Metadata=self.metadata.dump(),

cognite/neat/rules/models/asset/_validation.py

@@ -3,7 +3,7 @@ from typing import cast
 
 from cognite.neat.issues import IssueList
 from cognite.neat.issues.errors import NeatValueError, PropertyDefinitionError
-from cognite.neat.rules.models._base import SheetList
+from cognite.neat.rules.models._base_rules import SheetList
 from cognite.neat.rules.models.asset._rules import AssetProperty, AssetRules
 from cognite.neat.rules.models.entities import AssetEntity, AssetFields, ClassEntity
 from cognite.neat.rules.models.information._validation import InformationPostValidation

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",
 ]