Make AssetHierarchy work with AssetWrite (#1687)

CHANGELOG.md

@@ -17,7 +17,11 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
-## [7.32.5] - 2024-04-4
+## [7.32.6] - 2024-04-05
+### Fixed
+- `AssetsAPI.create_hierarchy` now properly supports `AssetWrite`.
+
+## [7.32.5] - 2024-04-04
 ### Improved
 - Type validation of identifiers
 

cognite/client/_api/assets.py

@@ -1252,7 +1252,7 @@ def _insert(
     ) -> _TaskResult:
         try:
             resp = self.assets_api._post(self.resource_path, self._dump_assets(assets))
-            successful = list(map(Asset.load, resp.json()["items"]))
+            successful = [Asset._load(item) for item in resp.json()["items"]]
             return _TaskResult(successful, failed=[], unknown=[])
         except Exception as err:
             self._set_latest_exception(err)

cognite/client/_version.py

@@ -1,4 +1,4 @@
 from __future__ import annotations
 
-__version__ = "7.32.5"
+__version__ = "7.32.6"
 __api_subversion__ = "20230101"

cognite/client/data_classes/_base.py

@@ -238,7 +238,7 @@ def _repr_html_(self) -> str:
 class WriteableCogniteResource(CogniteResource, Generic[T_WriteClass]):
     @abstractmethod
     def as_write(self) -> T_WriteClass:
-        raise NotImplementedError()
+        raise NotImplementedError
 
 
 T_CogniteResource = TypeVar("T_CogniteResource", bound=CogniteResource)
@@ -402,7 +402,7 @@ class WriteableCogniteResourceList(
 ):
     @abstractmethod
     def as_write(self) -> CogniteResourceList[T_WriteClass]:
-        raise NotImplementedError()
+        raise NotImplementedError
 
 
 @dataclass

cognite/client/data_classes/assets.py

@@ -724,20 +724,28 @@ class AssetHierarchy:
     """
 
     def __init__(self, assets: Sequence[Asset | AssetWrite], ignore_orphans: bool = False) -> None:
-        self._assets = assets
-        self._roots: list[Asset | AssetWrite] | None = None
-        self._orphans: list[Asset | AssetWrite] | None = None
+        self._assets = self._convert_to_read(assets)
+        self._roots: list[Asset] | None = None
+        self._orphans: list[Asset] | None = None
         self._ignore_orphans = ignore_orphans
-        self._invalid: list[Asset | AssetWrite] | None = None
-        self._unsure_parents: list[Asset | AssetWrite] | None = None
-        self._duplicates: dict[str, list[Asset | AssetWrite]] | None = None
+        self._invalid: list[Asset] | None = None
+        self._unsure_parents: list[Asset] | None = None
+        self._duplicates: dict[str, list[Asset]] | None = None
         self._cycles: list[list[str]] | None = None
 
         self.__validation_has_run = False
 
     def __len__(self) -> int:
         return len(self._assets)
 
+    @staticmethod
+    def _convert_to_read(assets: Sequence[Asset | AssetWrite]) -> Sequence[Asset]:
+        # TODO: AssetHierarchy doesn't work with AssetWrite (or more correctly, _AssetHierarchyCreator...)
+        #       and as we don't have a reverse of "as_write", we dump-load the write into a read:
+        return [
+            Asset._load(asset.dump(camel_case=True)) if isinstance(asset, AssetWrite) else asset for asset in assets
+        ]
+
     def is_valid(self, on_error: Literal["ignore", "warn", "raise"] = "ignore") -> bool:
         if not self.__validation_has_run:
             self.validate(verbose=False, on_error="ignore")
@@ -774,7 +782,7 @@ def unsure_parents(self) -> AssetList:
         return AssetList(self._unsure_parents)
 
     @property
-    def duplicates(self) -> dict[str, list[Asset | AssetWrite]]:
+    def duplicates(self) -> dict[str, list[Asset]]:
         if self._duplicates is None:
             raise RuntimeError("Unable to list duplicate assets before validation has run")
         # NB: Do not return AssetList (as it does not handle duplicates well):
@@ -814,7 +822,7 @@ def validate(
     def validate_and_report(self, output_file: Path | None = None) -> AssetHierarchy:
         return self.validate(verbose=True, output_file=output_file, on_error="ignore")
 
-    def groupby_parent_xid(self) -> dict[str | None, list[Asset | AssetWrite]]:
+    def groupby_parent_xid(self) -> dict[str | None, list[Asset]]:
         """Returns a mapping from parent external ID to a list of its direct children.
 
         Note:
@@ -824,14 +832,14 @@ def groupby_parent_xid(self) -> dict[str | None, list[Asset | AssetWrite]]:
             The same is true for all assets linking its parent by ID.
 
         Returns:
-            dict[str | None, list[Asset | AssetWrite]]: No description."""
+            dict[str | None, list[Asset]]: No description."""
         self.is_valid(on_error="raise")
 
         # Sort (on parent) as required by groupby. This is tricky as we need to avoid comparing string with None,
         # and can't simply hash it because of the possibility of collisions. Further, the empty string is a valid
         # external ID leaving no other choice than to prepend all strings with ' ' before comparison:
 
-        def parent_sort_fn(asset: Asset | AssetWrite) -> str:
+        def parent_sort_fn(asset: Asset) -> str:
             # All assets using 'parent_id' will be grouped together with the root assets:
             if (pxid := asset.parent_external_id) is None:
                 return ""
@@ -855,11 +863,11 @@ def parent_sort_fn(asset: Asset | AssetWrite) -> str:
         )
         return mapping
 
-    def count_subtree(self, mapping: dict[str | None, list[Asset | AssetWrite]]) -> dict[str, int]:
+    def count_subtree(self, mapping: dict[str | None, list[Asset]]) -> dict[str, int]:
         """Returns a mapping from asset external ID to the size of its subtree (children and children of children etc.).
 
         Args:
-            mapping (dict[str | None, list[Asset | AssetWrite]]): The mapping returned by `groupby_parent_xid()`. If None is passed, will be recreated (slightly expensive).
+            mapping (dict[str | None, list[Asset]]): The mapping returned by `groupby_parent_xid()`. If None is passed, will be recreated (slightly expensive).
 
         Returns:
             dict[str, int]: Lookup from external ID to descendant count.
@@ -894,23 +902,13 @@ def _no_basic_issues(self) -> bool:
             self._invalid or self._unsure_parents or self._duplicates or (self._orphans and not self._ignore_orphans)
         )
 
-    def _inspect_attributes(
-        self,
-    ) -> tuple[
-        list[Asset | AssetWrite],
-        list[Asset | AssetWrite],
-        list[Asset | AssetWrite],
-        list[Asset | AssetWrite],
-        dict[str, list[Asset | AssetWrite]],
-    ]:
-        invalid, orphans, roots, unsure_parents, duplicates = [], [], [], [], defaultdict(list)
+    def _inspect_attributes(self) -> tuple[list[Asset], list[Asset], list[Asset], list[Asset], dict[str, list[Asset]]]:
+        invalid, orphans, roots, unsure_parents = [], [], [], []  # type: tuple[list[Asset], list[Asset], list[Asset], list[Asset]]
+        duplicates: defaultdict[str, list[Asset]] = defaultdict(list)
         xid_count = Counter(a.external_id for a in self._assets)
 
         for asset in self._assets:
-            if isinstance(asset, AssetWrite):
-                id_, xid, name = None, asset.external_id, asset.name
-            else:
-                id_, xid, name = asset.id, asset.external_id, asset.name
+            id_, xid, name = asset.id, asset.external_id, asset.name
             if xid is None or name is None or len(name) < 1 or id_ is not None:
                 invalid.append(asset)
                 continue  # Don't report invalid as part of any other group
@@ -1005,7 +1003,7 @@ def print_header(title: str, columns: list[str]) -> None:
                 DrawTables.XLINE.join(DrawTables.HLINE * 20 for _ in columns),
             )
 
-        def print_table(lst: list[Asset | AssetWrite], columns: list[str]) -> None:
+        def print_table(lst: list[Asset], columns: list[str]) -> None:
             for entry in lst:
                 cols = (f"{shorten(getattr(entry, col)):<20}" for col in columns)
                 print_fn(DrawTables.VLINE.join(cols))

cognite/client/data_classes/data_modeling/instances.py

@@ -412,7 +412,7 @@ def to_pandas(  # type: ignore [override]
     @abstractmethod
     def as_apply(self) -> InstanceApply:
         """Convert the instance to an apply instance."""
-        raise NotImplementedError()
+        raise NotImplementedError
 
 
 class InstanceApplyResult(InstanceCore, ABC):
@@ -999,7 +999,7 @@ class InstancesResult:
 
     @classmethod
     def load(cls, data: str | dict) -> InstancesResult:
-        raise NotImplementedError()
+        raise NotImplementedError
 
 
 @dataclass

pyproject.toml

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

tests/tests_unit/test_data_classes/test_assets.py

@@ -15,6 +15,7 @@
     Asset,
     AssetHierarchy,
     AssetList,
+    AssetWrite,
     Event,
     EventList,
     FileMetadata,
@@ -152,15 +153,15 @@ def basic_issue_assets():
         Asset(name="a1", external_id="i am groot", parent_external_id=None),
         # Duplicated XIDs:
         Asset(name="a2", external_id="a", parent_external_id="i am groot"),
-        Asset(name="a3", external_id="a", parent_external_id="i am groot"),
+        AssetWrite(name="a3", external_id="a", parent_external_id="i am groot"),
         # Duplicated AND orphan:
         Asset(name="a4", external_id="a", parent_external_id="i am orphan"),
         # Orphan:
-        Asset(name="a5", external_id="b", parent_external_id="i am orphan"),
+        AssetWrite(name="a5", external_id="b", parent_external_id="i am orphan"),
         # Invalid (missing XIDs):
         Asset(name="a6", external_id=None, parent_external_id="i am groot"),
         # Doubly defined parent asset:
-        Asset(name="a7", external_id="c", parent_external_id="i am groot", parent_id=42),
+        AssetWrite(name="a7", external_id="c", parent_external_id="i am groot", parent_id=42),
     ]
 
 
@@ -237,9 +238,9 @@ class TestAssetHierarchy:
         (
             # Invalid name:
             Asset(name="", external_id="foo"),
-            Asset(name=None, external_id="foo"),
+            AssetWrite(name=None, external_id="foo"),
             # Invalid external_id (empty str allowed):
-            Asset(name="a", external_id=None),
+            AssetWrite(name="a", external_id=None),
             # Id given:
             Asset(name="a", external_id="", id=123),
         ),
@@ -253,7 +254,7 @@ def test_validate_asset_hierarchy___invalid_assets(self, asset):
     def test_validate_asset_hierarchy__orphans_given_ignore_false(self):
         assets = [
             Asset(name="a", parent_external_id="1", external_id="2"),
-            Asset(name="a", parent_external_id="2", external_id="3"),
+            AssetWrite(name="a", parent_external_id="2", external_id="3"),
         ]
         hierarchy = AssetHierarchy(assets, ignore_orphans=False).validate(on_error="ignore")
         assert len(hierarchy.orphans) == 1

tests/utils.py

@@ -201,7 +201,7 @@ def cdf_aggregate(
         raw_freq (str): The frequency of the raw data. If it is not given, it is attempted inferred from raw_df.
     """
     if is_step:
-        raise NotImplementedError()
+        raise NotImplementedError
 
     pd = cast(Any, local_import("pandas"))
     granularity_pd = granularity.replace("m", "T")