Enum property (#1851)

Co-authored-by: Håkon V. Treider <[email protected]>

CHANGELOG.md

@@ -17,6 +17,11 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## [7.54.3] - 2024-07-17
+### Added
+- [Feature Preview] Support for `enum` as container property type in the data modeling APIs. Note that this is not
+  yet supported in the API, and is an experimental feature that may change without warning.
+
 ## [7.54.2] - 2024-07-16
 ### Fixed
 - A bug in the list method of the RelationshipsAPI that could cause a thread deadlock.

cognite/client/_version.py

@@ -1,4 +1,4 @@
 from __future__ import annotations
 
-__version__ = "7.54.2"
+__version__ = "7.54.3"
 __api_subversion__ = "20230101"

cognite/client/data_classes/data_modeling/__init__.py

@@ -35,7 +35,6 @@
     Int32,
     Int64,
     Json,
-    Primitive,
     PropertyType,
     SequenceReference,
     Text,
@@ -152,7 +151,6 @@
     "Index",
     "Constraint",
     "ContainerProperty",
-    "Primitive",
     "CDFExternalIdReference",
     "RequiresConstraint",
     "ContainerId",

cognite/client/data_classes/data_modeling/data_types.py

@@ -3,7 +3,7 @@
 import logging
 from abc import ABC
 from dataclasses import asdict, dataclass
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar, cast
 
 from typing_extensions import Self, TypeAlias
 
@@ -52,7 +52,6 @@ def as_tuple(self) -> tuple[str, str]:
 @dataclass
 class PropertyType(CogniteObject, ABC):
     _type: ClassVar[str]
-    is_list: bool = False
 
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
         output = asdict(self)
@@ -71,78 +70,79 @@ def __load_unit_ref(data: dict) -> UnitReference | None:
         return unit
 
     @classmethod
-    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> PropertyType:
         type_ = resource["type"]
-        obj: Any
         if type_ == "text":
-            obj = Text(is_list=resource["list"], collation=resource.get("collation", "ucs_basic"))
+            return Text(is_list=resource["list"], collation=resource.get("collation", "ucs_basic"))
         elif type_ == "boolean":
-            obj = Boolean(is_list=resource["list"])
+            return Boolean(is_list=resource["list"])
         elif type_ == "float32":
-            obj = Float32(is_list=resource["list"], unit=cls.__load_unit_ref(resource))
+            return Float32(is_list=resource["list"], unit=cls.__load_unit_ref(resource))
         elif type_ == "float64":
-            obj = Float64(is_list=resource["list"], unit=cls.__load_unit_ref(resource))
+            return Float64(is_list=resource["list"], unit=cls.__load_unit_ref(resource))
         elif type_ == "int32":
-            obj = Int32(is_list=resource["list"], unit=cls.__load_unit_ref(resource))
+            return Int32(is_list=resource["list"], unit=cls.__load_unit_ref(resource))
         elif type_ == "int64":
-            obj = Int64(is_list=resource["list"], unit=cls.__load_unit_ref(resource))
+            return Int64(is_list=resource["list"], unit=cls.__load_unit_ref(resource))
         elif type_ == "timestamp":
-            obj = Timestamp(is_list=resource["list"])
+            return Timestamp(is_list=resource["list"])
         elif type_ == "date":
-            obj = Date(is_list=resource["list"])
+            return Date(is_list=resource["list"])
         elif type_ == "json":
-            obj = Json(is_list=resource["list"])
+            return Json(is_list=resource["list"])
         elif type_ == "timeseries":
-            obj = TimeSeriesReference(is_list=resource["list"])
+            return TimeSeriesReference(is_list=resource["list"])
         elif type_ == "file":
-            obj = FileReference(is_list=resource["list"])
+            return FileReference(is_list=resource["list"])
         elif type_ == "sequence":
-            obj = SequenceReference(is_list=resource["list"])
+            return SequenceReference(is_list=resource["list"])
         elif type_ == "direct":
-            obj = DirectRelation(
+            return DirectRelation(
                 container=ContainerId.load(container) if (container := resource.get("container")) else None,
                 # The PropertyTypes are used as both read and write objects. The `list` was added later
                 # in the API for DirectRelations. Thus, we need to set the default value to False
                 # to avoid breaking changes. When used as a read object, the `list` will always be present.
                 is_list=resource.get("list", False),
             )
-        else:
-            logger.warning(f"Unknown property type: {type_}")
-            obj = UnknownCogniteObject(resource)
-        return obj
+        elif type_ == "enum":
+            values = {key: EnumValue._load(value) for key, value in resource["values"].items()}
+            return Enum(values=values, unknown_value=resource.get("unknownValue"))
+        logger.warning(f"Unknown property type: {type_}")
+        return cast(Self, UnknownCogniteObject(resource))
 
 
 # Kept around for backwards compatibility
 UnknownPropertyType: TypeAlias = UnknownCogniteObject
 
 
 @dataclass
-class Text(PropertyType):
-    _type = "text"
-    collation: str = "ucs_basic"
+class ListablePropertyType(PropertyType, ABC):
+    is_list: bool = False
 
 
 @dataclass
-class Primitive(PropertyType, ABC): ...
+class Text(ListablePropertyType):
+    _type = "text"
+    collation: str = "ucs_basic"
 
 
 @dataclass
-class Boolean(PropertyType):
+class Boolean(ListablePropertyType):
     _type = "boolean"
 
 
 @dataclass
-class Timestamp(PropertyType):
+class Timestamp(ListablePropertyType):
     _type = "timestamp"
 
 
 @dataclass
-class Date(PropertyType):
+class Date(ListablePropertyType):
     _type = "date"
 
 
 @dataclass
-class Json(PropertyType):
+class Json(ListablePropertyType):
     _type = "json"
 
 
@@ -178,7 +178,7 @@ def load(cls, data: dict) -> UnitSystemReference:
 
 
 @dataclass
-class PropertyTypeWithUnit(PropertyType, ABC):
+class PropertyTypeWithUnit(ListablePropertyType, ABC):
     unit: UnitReference | None = None
 
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
@@ -209,7 +209,7 @@ class Int64(PropertyTypeWithUnit):
 
 
 @dataclass
-class CDFExternalIdReference(PropertyType, ABC): ...
+class CDFExternalIdReference(ListablePropertyType, ABC): ...
 
 
 @dataclass
@@ -228,7 +228,7 @@ class SequenceReference(CDFExternalIdReference):
 
 
 @dataclass
-class DirectRelation(PropertyType):
+class DirectRelation(ListablePropertyType):
     _type = "direct"
     container: ContainerId | None = None
 
@@ -240,3 +240,16 @@ def dump(self, camel_case: bool = True) -> dict:
             elif output["container"] is None:
                 output.pop("container")
         return output
+
+
+@dataclass
+class EnumValue(CogniteObject):
+    name: str | None = None
+    description: str | None = None
+
+
+@dataclass
+class Enum(PropertyType):
+    _type = "enum"
+    values: dict[str, EnumValue]
+    unknown_value: str | None = None

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "cognite-sdk"
 
-version = "7.54.2"
+version = "7.54.3"
 description = "Cognite Python SDK"
 readme = "README.md"
 documentation = "https://cognite-sdk-python.readthedocs-hosted.com"

tests/tests_unit/test_data_classes/test_data_models/test_data_types.py

@@ -1,12 +1,12 @@
 import pytest
 
+from cognite.client.data_classes._base import UnknownCogniteObject
 from cognite.client.data_classes.data_modeling.data_types import (
     DirectRelationReference,
     PropertyType,
     PropertyTypeWithUnit,
     Text,
     UnitReference,
-    UnknownPropertyType,
 )
 
 
@@ -37,12 +37,28 @@ class TestPropertyType:
                 "container": {"space": "mySpace", "externalId": "myId", "type": "container"},
                 "list": True,
             },
+            {
+                "type": "enum",
+                "values": {
+                    "string": {
+                        "name": "string",
+                        "description": "Time series with string data points.",
+                    },
+                    "numeric": {
+                        "name": "numeric",
+                        "description": "Time series with double floating point data points.",
+                    },
+                },
+            },
         ],
+        ids=lambda d: d.get("type", "unknown"),
     )
     def test_load_dump(self, data: dict) -> None:
-        actual = PropertyType.load(data).dump(camel_case=True)
+        loaded = PropertyType.load(data)
+        dumped = loaded.dump(camel_case=True)
 
-        assert data == actual
+        assert not isinstance(loaded, UnknownCogniteObject)
+        assert data == dumped
 
     def test_load_ignore_unknown_properties(self) -> None:
         data = {"type": "float64", "list": True, "unit": {"externalId": "known"}, "unknownProperty": "unknown"}
@@ -53,7 +69,7 @@ def test_load_ignore_unknown_properties(self) -> None:
     def test_load_dump_unkown_property(self) -> None:
         data = {"type": "unknowngibberish", "list": True, "unit": {"externalId": "known"}}
         obj = PropertyType.load(data)
-        assert isinstance(obj, UnknownPropertyType)
+        assert isinstance(obj, UnknownCogniteObject)
         actual = obj.dump(camel_case=True)
         assert data == actual