cognite-neat 0.123.26__tar.gz → 0.123.28__tar.gz

{cognite_neat-0.123.26 → cognite_neat-0.123.28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite-neat
-Version: 0.123.26
+Version: 0.123.28
 Summary: Knowledge graph transformation
 Project-URL: Documentation, https://cognite-neat.readthedocs-hosted.com/
 Project-URL: Homepage, https://cognite-neat.readthedocs-hosted.com/
@@ -18,7 +18,7 @@ Requires-Dist: jsonpath-python<2.0.0,>=1.0.6
 Requires-Dist: mixpanel<5.0.0,>=4.10.1
 Requires-Dist: networkx<4.0.0,>=3.4.2
 Requires-Dist: openpyxl<4.0.0,>=3.0.10
-Requires-Dist: packaging<25.0,>=22.0
+Requires-Dist: packaging>=22.0
 Requires-Dist: pandas<3.0.0,>=1.5.3
 Requires-Dist: pydantic<3.0.0,>=2.0.0
 Requires-Dist: pyvis<1.0.0,>=0.3.2

cognite_neat-0.123.28/cognite/neat/_version.py

@@ -0,0 +1,2 @@
+__version__ = "0.123.28"
+__engine__ = "^2.0.4"

{cognite_neat-0.123.26 → cognite_neat-0.123.28}/cognite/neat/core/_data_model/_constants.py

@@ -43,6 +43,10 @@ class EntityTypes(StrEnum):
     prefix = "prefix"
     space = "space"
     container_index = "container_index"
+    concept_restriction = "conceptRestriction"
+    value_constraint = "valueConstraint"
+    cardinality_constraint = "cardinalityConstraint"
+    named_individual = "named_individual"
 
 
 def get_reserved_words(

{cognite_neat-0.123.26 → cognite_neat-0.123.28}/cognite/neat/core/_data_model/models/conceptual/_unverified.py

@@ -100,8 +100,8 @@ class UnverifiedConceptualProperty(UnverifiedComponent[ConceptualProperty]):
 
     def dump(self, default_prefix: str, **kwargs) -> dict[str, Any]:  # type: ignore
         output = super().dump()
-        output["Concept"] = ConceptEntity.load(self.concept, prefix=default_prefix)
-        output["Value Type"] = load_value_type(self.value_type, default_prefix)
+        output["Concept"] = ConceptEntity.load(self.concept, prefix=default_prefix, return_on_failure=True)
+        output["Value Type"] = load_value_type(self.value_type, default_prefix, return_on_failure=True)
         return output
 
     def copy(self, update: dict[str, Any], default_prefix: str) -> "UnverifiedConceptualProperty":
@@ -135,10 +135,16 @@ class UnverifiedConcept(UnverifiedComponent[Concept]):
         parent: list[ConceptEntity] | None = None
         if isinstance(self.implements, str):
             self.implements = self.implements.strip()
-            parent = [ConceptEntity.load(parent, prefix=default_prefix) for parent in self.implements.split(",")]
+            parent = [
+                ConceptEntity.load(parent_str, prefix=default_prefix, return_on_failure=True)
+                for parent_str in self.implements.split(",")
+            ]
         elif isinstance(self.implements, list):
-            parent = [ConceptEntity.load(parent_, prefix=default_prefix) for parent_ in self.implements]
-        output["Concept"] = ConceptEntity.load(self.concept, prefix=default_prefix)
+            parent = [
+                ConceptEntity.load(parent_str, prefix=default_prefix, return_on_failure=True)
+                for parent_str in self.implements
+            ]
+        output["Concept"] = ConceptEntity.load(self.concept, prefix=default_prefix, return_on_failure=True)
         output["Implements"] = parent
         return output
 

{cognite_neat-0.123.26 → cognite_neat-0.123.28}/cognite/neat/core/_data_model/models/entities/__init__.py

@@ -1,6 +1,7 @@
 from ._constants import Undefined, Unknown
 from ._loaders import load_connection, load_dms_value_type, load_value_type
 from ._multi_value import MultiValueTypeInfo
+from ._restrictions import ConceptPropertyCardinalityConstraint, ConceptPropertyValueConstraint, parse_restriction
 from ._single_value import (
     AssetEntity,
     AssetFields,
@@ -31,6 +32,8 @@ __all__ = [
     "CdfResourceEntityList",
     "ClassEntityList",
     "ConceptEntity",
+    "ConceptPropertyCardinalityConstraint",
+    "ConceptPropertyValueConstraint",
     "ConceptualEntity",
     "ContainerEntity",
     "ContainerEntityList",
@@ -61,4 +64,5 @@ __all__ = [
     "load_connection",
     "load_dms_value_type",
     "load_value_type",
+    "parse_restriction",
 ]

cognite_neat-0.123.28/cognite/neat/core/_data_model/models/entities/_loaders.py

@@ -0,0 +1,155 @@
+from typing import Literal, overload
+
+from cognite.neat.core._data_model.models.data_types import DataType
+from cognite.neat.core._issues.errors import NeatTypeError
+
+from ._multi_value import MultiValueTypeInfo
+from ._single_value import (
+    ConceptEntity,
+    EdgeEntity,
+    PhysicalUnknownEntity,
+    ReverseConnectionEntity,
+    Unknown,
+    UnknownEntity,
+    ViewEntity,
+)
+
+
+@overload
+def load_value_type(
+    raw: str | MultiValueTypeInfo | DataType | ConceptEntity | UnknownEntity,
+    default_prefix: str,
+    return_on_failure: Literal[False] = False,
+) -> MultiValueTypeInfo | DataType | ConceptEntity | UnknownEntity: ...
+
+
+@overload
+def load_value_type(
+    raw: str | MultiValueTypeInfo | DataType | ConceptEntity | UnknownEntity,
+    default_prefix: str,
+    return_on_failure: Literal[True],
+) -> MultiValueTypeInfo | DataType | ConceptEntity | UnknownEntity | None | str: ...
+
+
+def load_value_type(
+    raw: str | MultiValueTypeInfo | DataType | ConceptEntity | UnknownEntity,
+    default_prefix: str,
+    return_on_failure: Literal[True, False] = False,
+) -> MultiValueTypeInfo | DataType | ConceptEntity | UnknownEntity | None | str:
+    """
+    Loads a value type from a raw string or entity.
+
+    Args:
+        raw: The raw value to load.
+        default_prefix: The default prefix to use if not specified in the raw value.
+        return_on_failure: If True, returns the raw value on parsing failure instead of raising an error.
+
+    Returns:
+        The loaded value type entity, or the raw value if loading fails and `return_on_failure` is True.
+    """
+    if isinstance(raw, MultiValueTypeInfo | DataType | ConceptEntity | UnknownEntity):
+        return raw
+    elif isinstance(raw, str):
+        # property holding xsd data type
+        # check if it is multi value type
+        if "," in raw:
+            value_type = MultiValueTypeInfo.load(raw)
+            value_type.set_default_prefix(default_prefix)
+            return value_type
+        elif DataType.is_data_type(raw):
+            return DataType.load(raw)
+
+        # unknown value type
+        elif raw == str(Unknown):
+            return UnknownEntity()
+
+        # property holding link to class
+        else:
+            return ConceptEntity.load(raw, prefix=default_prefix, return_on_failure=return_on_failure)
+    else:
+        raise NeatTypeError(f"Invalid value type: {type(raw)}")
+
+
+@overload
+def load_dms_value_type(
+    raw: str | DataType | ViewEntity | PhysicalUnknownEntity,
+    default_space: str,
+    default_version: str,
+    return_on_failure: Literal[False],
+) -> DataType | ViewEntity | PhysicalUnknownEntity: ...
+
+
+@overload
+def load_dms_value_type(
+    raw: str | DataType | ViewEntity | PhysicalUnknownEntity,
+    default_space: str,
+    default_version: str,
+    return_on_failure: Literal[True],
+) -> DataType | ViewEntity | PhysicalUnknownEntity | str: ...
+
+
+def load_dms_value_type(
+    raw: str | DataType | ViewEntity | PhysicalUnknownEntity,
+    default_space: str,
+    default_version: str,
+    return_on_failure: Literal[True, False] = False,
+) -> DataType | ViewEntity | PhysicalUnknownEntity | str:
+    """
+    Loads a value type from a raw string or entity in the context of a data modeling service
+
+    Args:
+        raw: The raw value to load.
+        default_space: The default space to use if not specified in the raw value.
+        default_version: The default version to use if not specified in the raw value.
+        return_on_failure: If True, returns the raw value on parsing failure instead of raising an error.
+
+    Returns:
+        The loaded value type entity, or the raw value if loading fails and `return_on_failure` is True.
+    """
+    if isinstance(raw, DataType | ViewEntity | PhysicalUnknownEntity):
+        return raw
+    elif isinstance(raw, str):
+        if DataType.is_data_type(raw):
+            return DataType.load(raw)
+        elif raw == str(Unknown):
+            return PhysicalUnknownEntity()
+        else:
+            return ViewEntity.load(
+                raw, space=default_space, version=default_version, return_on_failure=return_on_failure
+            )
+    raise NeatTypeError(f"Invalid value type: {type(raw)}")
+
+
+@overload
+def load_connection(
+    raw: Literal["direct"] | ReverseConnectionEntity | EdgeEntity | str | None,
+    default_space: str,
+    default_version: str,
+    return_on_failure: Literal[False] = False,
+) -> Literal["direct"] | ReverseConnectionEntity | EdgeEntity | None: ...
+
+
+@overload
+def load_connection(
+    raw: Literal["direct"] | ReverseConnectionEntity | EdgeEntity | str | None,
+    default_space: str,
+    default_version: str,
+    return_on_failure: Literal[True],
+) -> Literal["direct"] | ReverseConnectionEntity | EdgeEntity | None | str: ...
+
+
+def load_connection(
+    raw: Literal["direct"] | ReverseConnectionEntity | EdgeEntity | str | None,
+    default_space: str,
+    default_version: str,
+    return_on_failure: Literal[True, False] = False,
+) -> Literal["direct"] | ReverseConnectionEntity | EdgeEntity | None | str:
+    if isinstance(raw, str) and raw.lower() == "direct":
+        return "direct"  # type: ignore[return-value]
+    elif isinstance(raw, EdgeEntity | ReverseConnectionEntity) or raw is None:
+        return raw  # type: ignore[return-value]
+    elif isinstance(raw, str) and raw.startswith("edge"):
+        return EdgeEntity.load(raw, space=default_space, version=default_version, return_on_failure=return_on_failure)  # type: ignore[return-value]
+    elif isinstance(raw, str) and raw.startswith("reverse"):
+        return ReverseConnectionEntity.load(raw, return_on_failure=return_on_failure)  # type: ignore[return-value]
+    raise NeatTypeError(f"Invalid connection: {type(raw)}")

cognite_neat-0.123.28/cognite/neat/core/_data_model/models/entities/_restrictions.py

@@ -0,0 +1,230 @@
+import re
+import sys
+from abc import ABC
+from typing import Any, ClassVar, Final, Literal, TypeVar, cast, get_args
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from rdflib import Literal as RDFLiteral
+from rdflib import URIRef
+
+from cognite.neat.core._data_model._constants import EntityTypes
+from cognite.neat.core._data_model.models.data_types import _XSD_TYPES, DataType
+from cognite.neat.core._data_model.models.entities._constants import _PARSE
+from cognite.neat.core._data_model.models.entities._single_value import ConceptEntity, ConceptualEntity
+from cognite.neat.core._issues.errors._general import NeatValueError
+from cognite.neat.core._utils.rdf_ import remove_namespace_from_uri
+
+if sys.version_info <= (3, 11):
+    from typing_extensions import Self
+else:
+    from typing import Self
+
+
+ValueConstraints = Literal["allValuesFrom", "someValuesFrom", "hasValue"]
+CardinalityConstraints = Literal["minCardinality", "maxCardinality", "cardinality", "qualifiedCardinality"]
+
+
+# Constants for regex patterns - more maintainable
+PROPERTY_PATTERN: Final[str] = r"[a-zA-Z0-9._~?@!$&'*+,;=%-]+"
+VALUE_PATTERN: Final[str] = r".+"
+CARDINALITY_VALUE_PATTERN: Final[str] = r"\d+"
+ON_PATTERN: Final[str] = r"[^(]*"
+
+VALUE_CONSTRAINT_REGEX = re.compile(
+    rf"^{EntityTypes.value_constraint}:(?P<property>{PROPERTY_PATTERN})\((?P<constraint>{'|'.join(get_args(ValueConstraints))}),(?P<value>{VALUE_PATTERN})\)$"
+)
+
+CARDINALITY_CONSTRAINT_REGEX = re.compile(
+    rf"^{EntityTypes.cardinality_constraint}:(?P<property>{PROPERTY_PATTERN})\((?P<constraint>{'|'.join(get_args(CardinalityConstraints))}),(?P<value>{CARDINALITY_VALUE_PATTERN})(?:,(?P<on>{ON_PATTERN}))?\)$"
+)
+
+
+class NamedIndividualEntity(ConceptualEntity):
+    type_: ClassVar[EntityTypes] = EntityTypes.named_individual
+
+    @model_validator(mode="after")
+    def reset_prefix(self) -> Self:
+        self.prefix = "ni"
+        return self
+
+
+class ConceptPropertyRestriction(ABC, BaseModel):
+    model_config: ClassVar[ConfigDict] = ConfigDict(
+        str_strip_whitespace=True,
+        arbitrary_types_allowed=True,
+        strict=False,
+        extra="ignore",
+    )
+    type_: ClassVar[EntityTypes] = EntityTypes.concept_restriction
+    property_: str
+
+    @classmethod
+    def load(cls: "type[T_ConceptPropertyRestriction]", data: Any, **defaults: Any) -> "T_ConceptPropertyRestriction":
+        if isinstance(data, cls):
+            return data
+        elif not isinstance(data, str):
+            raise ValueError(f"Cannot load {cls.__name__} from {data}")
+
+        return cls.model_validate({_PARSE: data, "defaults": defaults})
+
+    @model_validator(mode="before")
+    def _load(cls, data: Any) -> dict:
+        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
+
+        if not isinstance(data, str):
+            raise ValueError(f"Cannot load {cls.__name__} from {data}")
+
+        return cls._parse(data, defaults)
+
+    @classmethod
+    def _parse(cls, data: str, defaults: dict) -> dict:
+        raise NotImplementedError(f"{cls.__name__} must implement _parse method")
+
+    def dump(self) -> str:
+        return self.__str__()
+
+    def as_tuple(self) -> tuple[str, ...]:
+        # We haver overwritten the serialization to str, so we need to do it manually
+        extra: tuple[str, ...] = tuple(
+            [
+                str(v or "")
+                for field_name in self.model_fields.keys()
+                if (v := getattr(self, field_name)) and field_name not in {"property_"}
+            ]
+        )
+
+        return self.property_, *extra
+
+    def __lt__(self, other: object) -> bool:
+        if not isinstance(other, ConceptPropertyRestriction):
+            return NotImplemented
+        return self.as_tuple() < other.as_tuple()
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, ConceptPropertyRestriction):
+            return NotImplemented
+        return self.as_tuple() == other.as_tuple()
+
+    def __hash__(self) -> int:
+        return hash(str(self))
+
+
+T_ConceptPropertyRestriction = TypeVar("T_ConceptPropertyRestriction", bound="ConceptPropertyRestriction")
+
+
+class ConceptPropertyValueConstraint(ConceptPropertyRestriction):
+    type_: ClassVar[EntityTypes] = EntityTypes.value_constraint
+    constraint: ValueConstraints
+    value: RDFLiteral | ConceptEntity | NamedIndividualEntity
+
+    @field_validator("value")
+    def validate_value(cls, value: Any) -> Any:
+        if isinstance(value, RDFLiteral) and value.datatype is None:
+            raise NeatValueError("RDFLiteral must have a datatype set, which must be one of the XSD types.")
+        return value
+
+    def __str__(self) -> str:
+        value_str = (
+            f"{self.value.value}^^{remove_namespace_from_uri(cast(URIRef, self.value.datatype))}"
+            if isinstance(self.value, RDFLiteral)
+            else str(self.value)
+        )
+        return f"{self.type_}:{self.property_}({self.constraint},{value_str})"
+
+    @classmethod
+    def _parse(cls, data: str, defaults: dict) -> dict:
+        if not (result := VALUE_CONSTRAINT_REGEX.match(data)):
+            raise NeatValueError(f"Invalid value constraint format: {data}")
+
+        property_ = result.group("property")
+        constraint = result.group("constraint")
+        raw_value = result.group("value")
+
+        value: NamedIndividualEntity | RDFLiteral | ConceptEntity
+        # scenario 1: NamedIndividual as value restriction
+        if raw_value.startswith("ni:"):
+            value = NamedIndividualEntity.load(raw_value)
+        # scenario 2: Datatype as value restriction
+        elif "^^" in raw_value:
+            if len(value_components := raw_value.split("^^")) == 2 and value_components[1] in _XSD_TYPES:
+                value = RDFLiteral(value_components[0], datatype=DataType.load(value_components[1]).as_xml_uri_ref())
+            else:
+                raise NeatValueError(f"Invalid value format for datatype: {raw_value}")
+
+        # scenario 3: ConceptEntity as value restriction
+        else:
+            value = ConceptEntity.load(raw_value, **defaults)
+
+        return dict(property_=property_, constraint=constraint, value=value)
+
+
+class ConceptPropertyCardinalityConstraint(ConceptPropertyRestriction):
+    type_: ClassVar[EntityTypes] = EntityTypes.cardinality_constraint
+    constraint: CardinalityConstraints
+    value: int = Field(ge=0)
+    on: DataType | ConceptEntity | None = None
+
+    def __str__(self) -> str:
+        on_str = f",{self.on}" if self.on else ""
+        return f"{self.type_}:{self.property_}({self.constraint},{self.value}{on_str})"
+
+    @classmethod
+    def _parse(cls, data: str, defaults: dict) -> dict:
+        if not (result := CARDINALITY_CONSTRAINT_REGEX.match(data)):
+            raise NeatValueError(f"Invalid cardinality constraint format: {data}")
+
+        property_ = result.group("property")
+        constraint = result.group("constraint")
+        value = result.group("value")
+        on = result.group("on")
+        if on:
+            if on in _XSD_TYPES:
+                on = DataType.load(on)
+            else:
+                on = cast(ConceptEntity, ConceptEntity.load(on, **defaults))
+
+        return dict(property_=property_, constraint=constraint, value=value, on=on)
+
+
+def parse_restriction(data: str, **defaults: Any) -> ConceptPropertyRestriction:
+    """Parse a string to create either a value or cardinality restriction.
+
+    Args:
+        data: String representation of the restriction
+        **defaults: Default values to use when parsing
+
+    Returns:
+        Either a ConceptPropertyValueConstraint or ConceptPropertyCardinalityConstraint
+
+    Raises:
+        NeatValueError: If the string cannot be parsed as either restriction type
+    """
+    # Check for value constraint pattern first (more specific)
+    if VALUE_CONSTRAINT_REGEX.match(data):
+        try:
+            return ConceptPropertyValueConstraint.load(data, **defaults)
+        except Exception as e:
+            raise NeatValueError(f"Failed to parse value constraint: {data}") from e
+
+    # Check for cardinality constraint pattern
+    if CARDINALITY_CONSTRAINT_REGEX.match(data):
+        try:
+            return ConceptPropertyCardinalityConstraint.load(data, **defaults)
+        except Exception as e:
+            raise NeatValueError(f"Failed to parse cardinality constraint: {data}") from e
+
+    # If neither pattern matches, provide a clear error
+    raise NeatValueError(
+        f"Invalid restriction format: {data}. "
+        f"Expected format: '{EntityTypes.value_constraint}:property(constraint,value)' "
+        f"or '{EntityTypes.cardinality_constraint}:property(constraint,value[,on])'"
+    )
+
+
+T_ConceptRestriction = TypeVar("T_ConceptRestriction", bound=ConceptPropertyRestriction)

{cognite_neat-0.123.26 → cognite_neat-0.123.28}/cognite/neat/core/_data_model/models/entities/_single_value.py

@@ -59,7 +59,13 @@ class ConceptualEntity(BaseModel, extra="ignore"):
 
     @classmethod
     @overload
-    def load(cls: "type[T_Entity]", data: Any, strict: Literal[True], **defaults: Any) -> "T_Entity": ...
+    def load(
+        cls: "type[T_Entity]",
+        data: Any,
+        strict: Literal[True],
+        return_on_failure: Literal[False] = False,
+        **defaults: Any,
+    ) -> "T_Entity": ...
 
     @classmethod
     @overload
@@ -67,25 +73,59 @@ class ConceptualEntity(BaseModel, extra="ignore"):
         cls: "type[T_Entity]",
         data: Any,
         strict: Literal[False] = False,
+        return_on_failure: Literal[False] = False,
         **defaults: Any,
     ) -> "T_Entity | UnknownEntity": ...
 
     @classmethod
-    def load(cls: "type[T_Entity]", data: Any, strict: bool = False, **defaults: Any) -> "T_Entity | UnknownEntity":
+    @overload
+    def load(
+        cls: "type[T_Entity]",
+        data: Any,
+        strict: Literal[False] = False,
+        return_on_failure: Literal[True] = True,
+        **defaults: Any,
+    ) -> "T_Entity | Any": ...
+
+    @classmethod
+    @overload
+    def load(
+        cls: "type[T_Entity]",
+        data: Any,
+        strict: Literal[True] = True,
+        return_on_failure: Literal[True] = True,
+        **defaults: Any,
+    ) -> "T_Entity | Any": ...
+
+    @classmethod
+    def load(
+        cls: "type[T_Entity]", data: Any, strict: bool = False, return_on_failure: bool = False, **defaults: Any
+    ) -> "T_Entity | UnknownEntity | Any":
+        """Loads an entity from a string or dict representation.
+
+        Args:
+            data: The data to load the entity from. Can be a string or a dict.
+            strict: If True, will raise an error if the data is "unknown".
+            return_on_failure: If True, will return the input data if loading fails.
+                This is used when you want to defer error handling to a later stage.
+            defaults: Default values to use when loading the entity. These will be used if the corresponding
+        """
         if isinstance(data, cls):
             return data
         elif isinstance(data, str) and data == str(Unknown):
             if strict:
                 raise NeatValueError(f"Failed to load entity {data!s}")
             return UnknownEntity(prefix=Undefined, suffix=Unknown)
-        if defaults and isinstance(defaults, dict):
-            # This is a trick to pass in default values
-            try:
+        try:
+            if defaults and isinstance(defaults, dict):
+                # This is a trick to pass in default values
                 return cls.model_validate({_PARSE: data, "defaults": defaults})
-            except ValueError:
-                raise
-        else:
-            return cls.model_validate(data)
+            else:
+                return cls.model_validate(data)
+        except ValueError:
+            if return_on_failure:
+                return data
+            raise
 
     @model_validator(mode="before")
     def _load(cls, data: Any) -> "dict | ConceptualEntity":
@@ -342,7 +382,13 @@ class PhysicalEntity(ConceptualEntity, Generic[T_ID], ABC):
 
     @classmethod  # type: ignore[override]
     @overload
-    def load(cls: "type[T_DMSEntity]", data: Any, strict: Literal[True], **defaults: Any) -> "T_DMSEntity": ...
+    def load(
+        cls: "type[T_DMSEntity]",
+        data: Any,
+        strict: Literal[True],
+        return_on_failure: Literal[False] = False,
+        **defaults: Any,
+    ) -> "T_DMSEntity": ...
 
     @classmethod
     @overload
@@ -350,18 +396,56 @@ class PhysicalEntity(ConceptualEntity, Generic[T_ID], ABC):
         cls: "type[T_DMSEntity]",
         data: Any,
         strict: Literal[False] = False,
+        return_on_failure: Literal[False] = False,
         **defaults: Any,
     ) -> "T_DMSEntity | PhysicalUnknownEntity": ...
 
     @classmethod
+    @overload
     def load(
-        cls: "type[T_DMSEntity]", data: Any, strict: bool = False, **defaults: Any
-    ) -> "T_DMSEntity | PhysicalUnknownEntity":  # type: ignore
+        cls: "type[T_DMSEntity]",
+        data: Any,
+        strict: Literal[False] = False,
+        return_on_failure: Literal[True] = True,
+        **defaults: Any,
+    ) -> "T_DMSEntity | Any": ...
+
+    @classmethod
+    @overload
+    def load(
+        cls: "type[T_DMSEntity]",
+        data: Any,
+        strict: Literal[True] = True,
+        return_on_failure: Literal[True] = True,
+        **defaults: Any,
+    ) -> "T_DMSEntity | Any": ...
+
+    @classmethod
+    def load(
+        cls: "type[T_DMSEntity]",
+        data: Any,
+        strict: Literal[True, False] = False,
+        return_on_failure: Literal[True, False] = False,
+        **defaults: Any,
+    ) -> "T_DMSEntity | PhysicalUnknownEntity | Any":
+        """Loads a DMS entity from a string or dict representation.
+
+        Args:
+            data: The data to load the entity from. Can be a string or a dict.
+            strict: If True, will raise an error if the data is "unknown".
+            return_on_failure: If True, will return the input data if loading fails.
+                This is used when you want to defer error handling to a later stage.
+            defaults: Default values to use when loading the entity. These will be used if the corresponding
+                fields are missing in the input data.
+
+        """
         if isinstance(data, str) and data == str(Unknown):
             if strict:
                 raise NeatValueError(f"Failed to load entity {data!s}")
             return PhysicalUnknownEntity.from_id(None)
-        return cast(T_DMSEntity, super().load(data, **defaults))
+        if isinstance(data, UnknownEntity):
+            return PhysicalUnknownEntity.from_id(None)
+        return cast(T_DMSEntity, super().load(data, strict=False, return_on_failure=return_on_failure, **defaults))
 
     @property
     def space(self) -> str: