fix annot. and rename _load_instance to _load_typed_instance

cognitedata · Sep 10, 2024 · b7bd67d · b7bd67d
1 parent d31620d
commit b7bd67d
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/cognite/client/data_classes/data_modeling/typed_instances.py b/cognite/client/data_classes/data_modeling/typed_instances.py
@@ -219,7 +219,7 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
     def _load(cls, resource: dict, cognite_client: CogniteClient | None = None) -> Self:
         sources = resource.pop("sources", [])
         properties = sources[0].get("properties", {}) if sources else {}
-        return cast(Self, _load_instance(cls, resource, properties, cls._instance_properties))
+        return cast(Self, _load_typed_instance(cls, resource, properties, cls._instance_properties))
 
 
 class TypedNodeApply(TypedInstanceApply):
@@ -372,7 +372,7 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
         all_properties = resource.pop("properties", {})
         source = cls.get_source()
         properties = all_properties.get(source.space, {}).get(source.as_source_identifier(), {})
-        return cast(Self, _load_instance(cls, resource, properties, cls._instance_properties))
+        return cast(Self, _load_typed_instance(cls, resource, properties, cls._instance_properties))
 
 
 class TypedEdge(TypedInstance):
@@ -451,12 +451,15 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
         all_properties = resource.pop("properties", {})
         source = cls.get_source()
         properties = all_properties.get(source.space, {}).get(source.as_source_identifier(), {})
-        return cast(Self, _load_instance(cls, resource, properties, cls._instance_properties))
+        return cast(Self, _load_typed_instance(cls, resource, properties, cls._instance_properties))
 
 
-def _load_instance(
-    cls: type, resource: dict[str, Any], properties: dict[str, Any], instance_properties: frozenset[str]
-) -> dict[str, Any]:
+def _load_typed_instance(
+    cls: type[TypedInstance | TypedInstanceApply],
+    resource: dict[str, Any],
+    properties: dict[str, Any],
+    instance_properties: frozenset[str],
+) -> TypedInstance | TypedInstanceApply:
     args: dict[str, Any] = {}
     resource.pop("instanceType", None)
     signature = inspect.signature(cls.__init__)  # type: ignore[misc]