cognite-neat 0.75.7__py3-none-any.whl → 0.75.9__py3-none-any.whl

cognite/neat/rules/models/entities.py

@@ -1,12 +1,11 @@
 import re
 import sys
-import threading
 from abc import ABC, abstractmethod
 from functools import total_ordering
-from typing import Annotated, Any, ClassVar, Generic, TypeVar
+from typing import Annotated, Any, ClassVar, Generic, TypeVar, cast
 
 from cognite.client.data_classes.data_modeling.ids import ContainerId, DataModelId, PropertyId, ViewId
-from pydantic import BaseModel, BeforeValidator, Field, PlainSerializer, model_serializer, model_validator
+from pydantic import AnyHttpUrl, BaseModel, BeforeValidator, Field, PlainSerializer, model_serializer, model_validator
 
 if sys.version_info >= (3, 11):
     from enum import StrEnum
@@ -58,45 +57,69 @@ _PROPERTY_ID_REGEX = rf"\((?P<{EntityTypes.property_}>{_ENTITY_ID_REGEX})\)"
 _ENTITY_PATTERN = re.compile(r"^(?P<prefix>.*?):?(?P<suffix>[^(:]*)(\((?P<content>[^)]+)\))?$")
 
 
-class _Undefined(BaseModel):
-    ...
+class _UndefinedType(BaseModel): ...
 
 
-class _Unknown(BaseModel):
+class _UnknownType(BaseModel):
     def __str__(self) -> str:
         return "#N/A"
 
 
 # This is a trick to make Undefined and Unknown singletons
-Undefined = _Undefined()
-Unknown = _Unknown()
+Undefined = _UndefinedType()
+Unknown = _UnknownType()
+_PARSE = object()
 
 
 @total_ordering
-class Entity(BaseModel):
+class Entity(BaseModel, extra="ignore"):
     """Entity is a class or property in OWL/RDF sense."""
 
     type_: ClassVar[EntityTypes] = EntityTypes.undefined
-    prefix: str | _Undefined = Undefined
-    suffix: str | _Unknown
+    prefix: str | _UndefinedType = Undefined
+    suffix: str
 
     @classmethod
-    def load(cls, data: Any) -> Self:
-        return cls.model_validate(data)
+    def load(cls: "type[T_Entity]", data: Any, **defaults) -> "T_Entity | UnknownEntity":
+        if isinstance(data, cls):
+            return data
+        elif isinstance(data, str) and data == str(Unknown):
+            return UnknownEntity(prefix=Undefined, suffix=Unknown)
+        if defaults and isinstance(defaults, dict):
+            # This is trick to pass in default values
+            return cls.model_validate({_PARSE: data, "defaults": defaults})
+        else:
+            return cls.model_validate(data)
 
     @model_validator(mode="before")
-    def _load(cls, data: Any) -> dict:
-        if isinstance(data, cls):
-            return data.model_dump()
-        elif isinstance(data, dict):
+    def _load(cls, data: Any) -> "dict | Entity":
+        defaults = {}
+        if isinstance(data, dict) and _PARSE in data:
+            defaults = data.get("defaults", {})
+            data = data[_PARSE]
+        if isinstance(data, dict):
+            data.update(defaults)
             return data
         elif hasattr(data, "versioned_id"):
             # Todo: Remove. Is here for backwards compatibility
             data = data.versioned_id
         elif not isinstance(data, str):
             raise ValueError(f"Cannot load {cls.__name__} from {data}")
-
-        return cls._parse(data)
+        elif data == str(Unknown) and cls.type_ == EntityTypes.undefined:
+            return dict(prefix=Undefined, suffix=Unknown)  # type: ignore[arg-type]
+        elif data == str(Unknown):
+            raise ValueError(f"Unknown is not allowed for {cls.type_} entity")
+
+        result = cls._parse(data)
+        output = defaults.copy()
+        # Populate by alias
+        for field_name, field_ in cls.model_fields.items():
+            name = field_.alias or field_name
+            if (field_value := result.get(field_name)) and not (field_value in [Unknown, Undefined] and name in output):
+                output[name] = result.pop(field_name)
+            elif name not in output and name in result:
+                output[name] = result.pop(name)
+        return output
 
     @model_serializer(when_used="unless-none", return_type=str)
     def as_str(self) -> str:
@@ -131,7 +154,7 @@ class Entity(BaseModel):
                 if isinstance(v := getattr(self, field_name), str | None) and field_name not in {"prefix", "suffix"}
             ]
         )
-        if isinstance(self.prefix, _Undefined):
+        if isinstance(self.prefix, _UndefinedType):
             return str(self.suffix), *extra
         else:
             return self.prefix, str(self.suffix), *extra
@@ -188,10 +211,25 @@ class Entity(BaseModel):
         return f"{self.prefix}:{self.suffix!s}"
 
 
+T_Entity = TypeVar("T_Entity", bound=Entity)
+
+
 class ClassEntity(Entity):
     type_: ClassVar[EntityTypes] = EntityTypes.class_
     version: str | None = None
 
+    def as_view_entity(self, default_space: str, default_version) -> "ViewEntity":
+        if self.version is None:
+            version = default_version
+        else:
+            version = self.version
+        space = default_space if isinstance(self.prefix, _UndefinedType) else self.prefix
+        return ViewEntity(space=space, externalId=str(self.suffix), version=version)
+
+    def as_container_entity(self, default_space: str) -> "ContainerEntity":
+        space = default_space if isinstance(self.prefix, _UndefinedType) else self.prefix
+        return ContainerEntity(space=space, externalId=str(self.suffix))
+
 
 class ParentClassEntity(ClassEntity):
     type_: ClassVar[EntityTypes] = EntityTypes.parent_class
@@ -200,25 +238,34 @@ class ParentClassEntity(ClassEntity):
         return ClassEntity(prefix=self.prefix, suffix=self.suffix, version=self.version)
 
 
-T_ID = TypeVar("T_ID", bound=ContainerId | ViewId | DataModelId | PropertyId)
+class UnknownEntity(ClassEntity):
+    type_: ClassVar[EntityTypes] = EntityTypes.undefined
+    prefix: _UndefinedType = Undefined
+    suffix: _UnknownType = Unknown  # type: ignore[assignment]
+
+    @property
+    def id(self) -> str:
+        return str(Unknown)
+
+
+T_ID = TypeVar("T_ID", bound=ContainerId | ViewId | DataModelId | PropertyId | None)
 
 
 class DMSEntity(Entity, Generic[T_ID], ABC):
     type_: ClassVar[EntityTypes] = EntityTypes.undefined
-    default_space_by_thread: ClassVar[dict[threading.Thread, str]] = {}
-    suffix: str
+    prefix: str = Field(alias="space")
+    suffix: str = Field(alias="externalId")
 
     @classmethod
-    def set_default_space(cls, space: str) -> None:
-        cls.default_space_by_thread[threading.current_thread()] = space
+    def load(cls: "type[T_DMSEntity]", data: Any, **defaults) -> "T_DMSEntity | DMSUnknownEntity":  # type: ignore[override]
+        if isinstance(data, str) and data == str(Unknown):
+            return DMSUnknownEntity.from_id(None)
+        return cast(T_DMSEntity, super().load(data, **defaults))
 
     @property
     def space(self) -> str:
         """Returns entity space in CDF."""
-        if isinstance(self.prefix, _Undefined):
-            return self.default_space_by_thread.get(threading.current_thread(), "MISSING")
-        else:
-            return self.prefix
+        return self.prefix
 
     @property
     def external_id(self) -> str:
@@ -229,6 +276,17 @@ class DMSEntity(Entity, Generic[T_ID], ABC):
     def as_id(self) -> T_ID:
         raise NotImplementedError("Method as_id must be implemented in subclasses")
 
+    @classmethod
+    @abstractmethod
+    def from_id(cls, id: T_ID) -> Self:
+        raise NotImplementedError("Method from_id must be implemented in subclasses")
+
+    def as_class(self) -> ClassEntity:
+        return ClassEntity(prefix=self.space, suffix=self.external_id)
+
+
+T_DMSEntity = TypeVar("T_DMSEntity", bound=DMSEntity)
+
 
 class ContainerEntity(DMSEntity[ContainerId]):
     type_: ClassVar[EntityTypes] = EntityTypes.container
@@ -236,16 +294,16 @@ class ContainerEntity(DMSEntity[ContainerId]):
     def as_id(self) -> ContainerId:
         return ContainerId(space=self.space, external_id=self.external_id)
 
+    @classmethod
+    def from_id(cls, id: ContainerId) -> "ContainerEntity":
+        return cls(space=id.space, externalId=id.external_id)
+
 
 class DMSVersionedEntity(DMSEntity[T_ID], ABC):
-    version: str | None = None
-    default_version_by_thread: ClassVar[dict[threading.Thread, str]] = {}
+    version: str
 
-    @property
-    def version_with_fallback(self) -> str:
-        if self.version is not None:
-            return self.version
-        return self.default_version_by_thread.get(threading.current_thread(), "MISSING")
+    def as_class(self) -> ClassEntity:
+        return ClassEntity(prefix=self.space, suffix=self.external_id, version=self.version)
 
 
 class ViewEntity(DMSVersionedEntity[ViewId]):
@@ -254,16 +312,54 @@ class ViewEntity(DMSVersionedEntity[ViewId]):
     def as_id(
         self,
     ) -> ViewId:
-        return ViewId(space=self.space, external_id=self.external_id, version=self.version_with_fallback)
+        return ViewId(space=self.space, external_id=self.external_id, version=self.version)
+
+    @classmethod
+    def from_id(cls, id: ViewId) -> "ViewEntity":
+        if id.version is None:
+            raise ValueError("Version must be specified")
+        return cls(space=id.space, externalId=id.external_id, version=id.version)
+
+
+class DMSUnknownEntity(DMSEntity[None]):
+    """This is a special entity that represents an unknown entity.
+
+    The use case is for direct relations where the source is not known."""
+
+    type_: ClassVar[EntityTypes] = EntityTypes.undefined
+    prefix: _UndefinedType = Field(Undefined, alias="space")  # type: ignore[assignment]
+    suffix: _UnknownType = Field(Unknown, alias="externalId")  # type: ignore[assignment]
+
+    def as_id(self) -> None:
+        return None
 
+    @classmethod
+    def from_id(cls, id: None) -> "DMSUnknownEntity":
+        return cls(space=Undefined, externalId=Unknown)
 
-class PropertyEntity(DMSVersionedEntity[PropertyId]):
+    @property
+    def id(self) -> str:
+        return str(Unknown)
+
+
+class ViewPropertyEntity(DMSVersionedEntity[PropertyId]):
     type_: ClassVar[EntityTypes] = EntityTypes.property_
     property_: str = Field(alias="property")
 
     def as_id(self) -> PropertyId:
-        return PropertyId(
-            source=ViewId(self.space, self.external_id, self.version_with_fallback), property=self.property_
+        return PropertyId(source=ViewId(self.space, self.external_id, self.version), property=self.property_)
+
+    def as_view_id(self) -> ViewId:
+        return ViewId(space=self.space, external_id=self.external_id, version=self.version)
+
+    @classmethod
+    def from_id(cls, id: PropertyId) -> "ViewPropertyEntity":
+        if isinstance(id.source, ContainerId):
+            raise ValueError("Only view source are supported")
+        if id.source.version is None:
+            raise ValueError("Version must be specified")
+        return cls(
+            space=id.source.space, externalId=id.source.external_id, version=id.source.version, property=id.property
         )
 
 
@@ -271,11 +367,32 @@ class DataModelEntity(DMSVersionedEntity[DataModelId]):
     type_: ClassVar[EntityTypes] = EntityTypes.datamodel
 
     def as_id(self) -> DataModelId:
-        return DataModelId(space=self.space, external_id=self.external_id, version=self.version_with_fallback)
+        return DataModelId(space=self.space, external_id=self.external_id, version=self.version)
+
+    @classmethod
+    def from_id(cls, id: DataModelId) -> "DataModelEntity":
+        if id.version is None:
+            raise ValueError("Version must be specified")
+        return cls(space=id.space, externalId=id.external_id, version=id.version)
 
 
-class ReferenceEntity(PropertyEntity):
+class ReferenceEntity(ClassEntity):
     type_: ClassVar[EntityTypes] = EntityTypes.reference_entity
+    prefix: str
+    property_: str | None = Field(None, alias="property")
+
+    def as_view_id(self) -> ViewId:
+        if isinstance(self.prefix, _UndefinedType) or isinstance(self.suffix, _UnknownType):
+            raise ValueError("Prefix is not defined or suffix is unknown")
+        return ViewId(space=self.prefix, external_id=self.suffix, version=self.version)
+
+    def as_view_property_id(self) -> PropertyId:
+        if self.property_ is None or self.prefix is Undefined or self.suffix is Unknown:
+            raise ValueError("Property is not defined or prefix is not defined or suffix is unknown")
+        return PropertyId(source=self.as_view_id(), property=self.property_)
+
+    def as_class_entity(self) -> ClassEntity:
+        return ClassEntity(prefix=self.prefix, suffix=self.suffix, version=self.version)
 
 
 def _split_str(v: Any) -> list[str]:
@@ -297,3 +414,28 @@ ParentEntityList = Annotated[
         when_used="unless-none",
     ),
 ]
+
+ContainerEntityList = Annotated[
+    list[ContainerEntity],
+    BeforeValidator(_split_str),
+    PlainSerializer(
+        _join_str,
+        return_type=str,
+        when_used="unless-none",
+    ),
+]
+
+ViewEntityList = Annotated[
+    list[ViewEntity],
+    BeforeValidator(_split_str),
+    PlainSerializer(
+        _join_str,
+        return_type=str,
+        when_used="unless-none",
+    ),
+]
+
+URLEntity = Annotated[
+    AnyHttpUrl,
+    PlainSerializer(lambda v: str(v), return_type=str, when_used="unless-none"),
+]

cognite/neat/rules/models/rdfpath.py

@@ -1,25 +1,15 @@
-"""
-"""
+""" """
 
 import re
 import sys
 from collections import Counter
-from typing import Literal
+from functools import total_ordering
+from typing import ClassVar, Literal
 
 from pydantic import BaseModel, field_validator
 
 from cognite.neat.rules import exceptions
 
-from ._entity import (
-    CLASS_ID_REGEX,
-    CLASS_ID_REGEX_COMPILED,
-    ENTITY_ID_REGEX,
-    PROPERTY_ID_REGEX,
-    SUFFIX_REGEX,
-    Entity,
-    EntityTypes,
-)
-
 if sys.version_info >= (3, 11):
     from enum import StrEnum
     from typing import Self
@@ -40,6 +30,26 @@ class Lookup(StrEnum):
     value = "value"  # type: ignore
 
 
+class EntityTypes(StrEnum):
+    class_ = "class"
+    property_ = "property"
+    undefined = "undefined"
+
+
+# FOR PARSING STRINGS:
+PREFIX_REGEX = r"[a-zA-Z]+[a-zA-Z0-9-_.]*[a-zA-Z0-9]+"
+SUFFIX_REGEX = r"[a-zA-Z0-9-_.]+[a-zA-Z0-9]|[-_.]*[a-zA-Z0-9]+"
+VERSION_REGEX = r"[a-zA-Z0-9]([.a-zA-Z0-9_-]{0,41}[a-zA-Z0-9])?"
+
+ENTITY_ID_REGEX = rf"{PREFIX_REGEX}:({SUFFIX_REGEX})"
+ENTITY_ID_REGEX_COMPILED = re.compile(rf"^(?P<prefix>{PREFIX_REGEX}):(?P<suffix>{SUFFIX_REGEX})$")
+VERSIONED_ENTITY_REGEX_COMPILED = re.compile(
+    rf"^(?P<prefix>{PREFIX_REGEX}):(?P<suffix>{SUFFIX_REGEX})\(version=(?P<version>{VERSION_REGEX})\)$"
+)
+CLASS_ID_REGEX = rf"(?P<{EntityTypes.class_}>{ENTITY_ID_REGEX})"
+CLASS_ID_REGEX_COMPILED = re.compile(rf"^{CLASS_ID_REGEX}$")
+PROPERTY_ID_REGEX = rf"\((?P<{EntityTypes.property_}>{ENTITY_ID_REGEX})\)"
+
 # traversal direction
 DIRECTION_REGEX = r"(?P<direction>(->|<-))"
 
@@ -73,6 +83,95 @@ TABLE_REGEX_COMPILED = re.compile(
 StepDirection = Literal["source", "target", "origin"]
 _direction_by_symbol: dict[str, StepDirection] = {"->": "target", "<-": "source"}
 
+Undefined = type(object())
+Unknown = type(object())
+
+
+# mypy does not like the sentinel value, and it is not possible to ignore only the line with it below.
+# so we ignore all errors beyond this point.
+# mypy: ignore-errors
+@total_ordering
+class Entity(BaseModel, arbitrary_types_allowed=True):
+    """Entity is a class or property in OWL/RDF sense."""
+
+    type_: ClassVar[EntityTypes] = EntityTypes.undefined
+    prefix: str | Undefined = Undefined
+    suffix: str | Unknown
+    version: str | None = None
+    name: str | None = None
+    description: str | None = None
+
+    def __lt__(self, other: object) -> bool:
+        if type(self) is not type(other) or not isinstance(other, Entity):
+            return NotImplemented
+        return self.versioned_id < other.versioned_id
+
+    def __eq__(self, other: object) -> bool:
+        if type(self) is not type(other) or not isinstance(other, Entity):
+            return NotImplemented
+        return self.versioned_id == other.versioned_id
+
+    def __hash__(self) -> int:
+        return hash(self.versioned_id)
+
+    def as_non_versioned_entity(self) -> Self:
+        return self.from_string(f"{self.prefix}:{self.suffix}")
+
+    @property
+    def id(self) -> str:
+        if self.suffix is Unknown:
+            return "#N/A"
+        elif self.prefix is Undefined:
+            return self.suffix
+        else:
+            return f"{self.prefix}:{self.suffix}"
+
+    @property
+    def versioned_id(self) -> str:
+        if self.version is None:
+            return self.id
+        else:
+            return f"{self.id}(version={self.version})"
+
+    @property
+    def space(self) -> str:
+        """Returns entity space in CDF."""
+        return self.prefix
+
+    @property
+    def external_id(self) -> str:
+        """Returns entity external id in CDF."""
+        return self.suffix
+
+    def __repr__(self):
+        return self.versioned_id
+
+    def __str__(self):
+        return self.versioned_id
+
+    @classmethod
+    def from_string(cls, entity_string: str, base_prefix: str | None = None) -> Self:
+        if entity_string == "#N/A":
+            return cls(prefix=Undefined, suffix=Unknown)
+        elif result := VERSIONED_ENTITY_REGEX_COMPILED.match(entity_string):
+            return cls(
+                prefix=result.group("prefix"),
+                suffix=result.group("suffix"),
+                version=result.group("version"),
+            )
+        elif result := ENTITY_ID_REGEX_COMPILED.match(entity_string):
+            return cls(prefix=result.group("prefix"), suffix=result.group("suffix"))
+        elif base_prefix and re.match(SUFFIX_REGEX, entity_string) and re.match(PREFIX_REGEX, base_prefix):
+            return cls(prefix=base_prefix, suffix=entity_string)
+        else:
+            raise ValueError(f"{cls.__name__} is expected to be prefix:suffix, got {entity_string}")
+
+    @classmethod
+    def from_list(cls, entity_strings: list[str], base_prefix: str | None = None) -> list[Self]:
+        return [
+            cls.from_string(entity_string=entity_string, base_prefix=base_prefix) for entity_string in entity_strings
+        ]
+
 
 class Step(BaseModel):
     class_: Entity

cognite/neat/rules/models/rules/_base.py

@@ -27,7 +27,7 @@ from pydantic import (
 )
 from pydantic.fields import FieldInfo
 
-from cognite.neat.rules.models.rules._types import ClassType
+from cognite.neat.rules.models.entities import ClassEntity
 
 if sys.version_info >= (3, 11):
     from enum import StrEnum
@@ -274,7 +274,7 @@ class BaseRules(RuleModel):
 
 # An sheet entity is either a class or a property.
 class SheetEntity(RuleModel):
-    class_: ClassType = Field(alias="Class")
+    class_: ClassEntity = Field(alias="Class")
     name: str | None = Field(alias="Name", default=None)
     description: str | None = Field(alias="Description", default=None)