diff --git a/cognite/client/data_classes/_base.py b/cognite/client/data_classes/_base.py index 8274fdd29e..b2b518dbec 100644 --- a/cognite/client/data_classes/_base.py +++ b/cognite/client/data_classes/_base.py @@ -27,7 +27,7 @@ from typing_extensions import TypeAlias from cognite.client.exceptions import CogniteMissingClientError -from cognite.client.utils._auxiliary import fast_dict_load, json_dump_default +from cognite.client.utils._auxiliary import fast_dict_load, get_identifiers, json_dump_default from cognite.client.utils._identifier import IdentifierSequence from cognite.client.utils._importing import local_import from cognite.client.utils._pandas_helpers import convert_nullable_int_cols, notebook_display_with_fallback @@ -204,14 +204,9 @@ def __init__(self, resources: Collection[Any], cognite_client: CogniteClient | N ) self._cognite_client = cast("CogniteClient", cognite_client) super().__init__(resources) - self._id_to_item, self._external_id_to_item = {}, {} - if self.data: - if hasattr(self.data[0], "external_id"): - self._external_id_to_item = { - item.external_id: item for item in self.data if item.external_id is not None - } - if hasattr(self.data[0], "id"): - self._id_to_item = {item.id: item for item in self.data if item.id is not None} + + self._identifier_to_items: dict[tuple[str, str | int], T_CogniteResource] = {} + self._create_identifier_mappings() def pop(self, i: int = -1) -> T_CogniteResource: return super().pop(i) @@ -242,10 +237,9 @@ def __str__(self) -> str: # TODO: We inherit a lot from UserList that we don't actually support... def extend(self, other: Collection[Any]) -> None: # type: ignore [override] other_res_list = type(self)(other) # See if we can accept the types - if set(self._id_to_item).isdisjoint(other_res_list._id_to_item): + if set(self._identifier_to_items).isdisjoint(other_res_list._identifier_to_items): super().extend(other) - self._external_id_to_item.update(other_res_list._external_id_to_item) - self._id_to_item.update(other_res_list._id_to_item) + self._identifier_to_items.update(other_res_list._identifier_to_items) else: raise ValueError("Unable to extend as this would introduce duplicates") @@ -260,21 +254,6 @@ def dump(self, camel_case: bool = False) -> list[dict[str, Any]]: """ return [resource.dump(camel_case) for resource in self.data] - def get(self, id: int | None = None, external_id: str | None = None) -> T_CogniteResource | None: - """Get an item from this list by id or external_id. - - Args: - id (int | None): The id of the item to get. - external_id (str | None): The external_id of the item to get. - - Returns: - T_CogniteResource | None: The requested item - """ - IdentifierSequence.load(id, external_id).assert_singleton() - if id: - return self._id_to_item.get(id) - return self._external_id_to_item.get(external_id) - def to_pandas( self, camel_case: bool = False, @@ -681,47 +660,205 @@ def dump(self, camel_case: bool = False) -> dict[str, Any]: T_CogniteSort = TypeVar("T_CogniteSort", bound=CogniteSort) -class HasExternalAndInternalId(Protocol): +class HasInternalId(Protocol): + @property + def id(self) -> int | None: + ... + + +class HasExternalId(Protocol): @property def external_id(self) -> str | None: ... + +class HasExtAndInternalId(HasInternalId, HasExternalId, Protocol): + ... + + +class HasUserIdentifier(Protocol): @property - def id(self) -> int | None: + def user_identifier(self) -> str: + ... + + +class HasName(Protocol): + @property + def name(self) -> str | None: + ... + + +class HasKey(Protocol): + @property + def key(self) -> str | None: ... -class IdTransformerMixin(Sequence[HasExternalAndInternalId], ABC): - def as_external_ids(self) -> list[str]: +class IdTransformerMixin(Sequence[HasInternalId], ABC): + def _create_identifier_mappings(self) -> None: + self._identifier_to_items.update( + {("id", item.id): item for item in self.data if item.id is not None}, + ) + + def get(self, id: int) -> T_CogniteResource | None: + """Get an item from this list by 'id'. + + Args: + id (int): The id of the item to get. + Returns: + T_CogniteResource | None: The requested item """ - Returns the external ids of all resources. + return self._identifier_to_items.get(("id", id)) + def as_ids(self, ignore_missing: bool = False) -> list[int]: + """ + Returns the ids of all resources. + + Args: + ignore_missing (bool): Skip when resources are missing 'id' instead of raising. Raises: - ValueError: If any resource in the list does not have an external id. + ValueError: If any resource in the list does not have an 'id'. + Returns: + list[int]: The ids of all resources in the list. + """ + return get_identifiers(self, "id", ignore_missing) + +class ExtIdTransformerMixin(Sequence[HasExternalId], ABC): + def _create_identifier_mappings(self) -> None: + self._identifier_to_items.update( + {("external_id", item.external_id): item for item in self.data if item.external_id is not None}, + ) + + def get(self, external_id: str) -> T_CogniteResource | None: + """Get an item from this list by 'external_id'. + + Args: + external_id (str): The external_id of the item to get. + Returns: + T_CogniteResource | None: The requested item + """ + return self._identifier_to_items.get(("external_id", external_id)) + + def as_external_ids(self, ignore_missing: bool = False) -> list[str]: + """ + Returns the external ids of all resources. + + Args: + ignore_missing (bool): Skip when resources are missing 'external id' instead of raising. + Raises: + ValueError: If any resource in the list does not have an 'external id'. Returns: list[str]: The external ids of all resources in the list. """ - external_ids: list[str] = [] - for x in self: - if x.external_id is None: - raise ValueError(f"All {type(x).__name__} must have external_id") - external_ids.append(x.external_id) - return external_ids - - def as_ids(self) -> list[int]: + return get_identifiers(self, "external_id", ignore_missing) + + +class IdAndExtIdTransformerMixin(IdTransformerMixin, ExtIdTransformerMixin, ABC): + def _create_identifier_mappings(self) -> None: + super()._create_identifier_mappings() # updates ids + super(IdTransformerMixin, self)._create_identifier_mappings() # updates external ids (yes) + + def get(self, id: int | None = None, external_id: str | None = None) -> T_CogniteResource | None: + """Get an item from this list by 'id' or 'external_id'. + + Args: + id (int | None): The id of the item to get. + external_id (str | None): The external_id of the item to get. + Returns: + T_CogniteResource | None: The requested item """ - Returns the ids of all resources. + IdentifierSequence.load(id, external_id).assert_singleton() + if id: + return self._identifier_to_items.get(("id", id)) + return self._identifier_to_items.get(("external_id", external_id)) + +class NameTransformerMixin(Sequence[HasName], ABC): + def _create_identifier_mappings(self) -> None: + self._identifier_to_items.update( + {("name", item.name): item for item in self.data if item.name is not None}, + ) + + def get(self, name: str) -> T_CogniteResource | None: + """Get an item from this list by 'name'. + + Args: + name (str): The name of the item to get. + Returns: + T_CogniteResource | None: The requested item + """ + return self._identifier_to_items.get(("name", name)) + + def as_names(self, ignore_missing: bool = False) -> list[str]: + """ + Returns the name of all resources. + + Args: + ignore_missing (bool): Skip when resources are missing 'name' instead of raising. Raises: - ValueError: If any resource in the list does not have an id. + ValueError: If any resource in the list does not have a 'name'. + Returns: + list[str]: The names of all resources in the list. + """ + return get_identifiers(self, "name", ignore_missing) + +class KeyTransformerMixin(Sequence[HasKey], ABC): + def _create_identifier_mappings(self) -> None: + self._identifier_to_items.update( + {("key", item.key): item for item in self.data if item.key is not None}, + ) + + def get(self, key: str) -> T_CogniteResource | None: + """Get an item from this list by 'key'. + + Args: + key (str): The key of the item to get. Returns: - list[int]: The ids of all resources in the list. + T_CogniteResource | None: The requested item + """ + return self._identifier_to_items.get(("key", key)) + + def as_keys(self, ignore_missing: bool = False) -> list[str]: + """ + Returns the keys of all resources. + + Args: + ignore_missing (bool): Skip when resources are missing 'key' instead of raising. + Raises: + ValueError: If any resource in the list does not have a 'key'. + Returns: + list[str]: The keys of all resources in the list. + """ + return get_identifiers(self, "key", ignore_missing) + + +class UserIdentTransformerMixin(Sequence[HasUserIdentifier], ABC): + def _create_identifier_mappings(self) -> None: + self._identifier_to_items.update( + {("user_identifier", item.user_identifier): item for item in self.data if item.user_identifier is not None} + ) + + def get(self, user_identifier: str) -> T_CogniteResource | None: + """Get an item from this list by 'user_identifier'. + + Args: + user_identifier (str): The user_identifier of the item to get. + Returns: + T_CogniteResource | None: The requested item + """ + return self._identifier_to_items.get(("user_identifier", user_identifier)) + + def as_user_identifiers(self, ignore_missing: bool = False) -> list[str]: + """ + Returns the user identifiers of all resources. + + Args: + ignore_missing (bool): Skip when resources are missing 'user identifier' instead of raising. + Raises: + ValueError: If any resource in the list does not have a 'user identifier'. + Returns: + list[str]: The user identifiers of all resources in the list. """ - ids: list[int] = [] - for x in self: - if x.id is None: - raise ValueError(f"All {type(x).__name__} must have id") - ids.append(x.id) - return ids + return get_identifiers(self, "user_identifier", ignore_missing) diff --git a/cognite/client/data_classes/assets.py b/cognite/client/data_classes/assets.py index 95cb22c2f4..ffa7a2dda6 100644 --- a/cognite/client/data_classes/assets.py +++ b/cognite/client/data_classes/assets.py @@ -38,7 +38,7 @@ CogniteSort, CogniteUpdate, EnumProperty, - IdTransformerMixin, + IdAndExtIdTransformerMixin, NoCaseConversionPropertyList, PropertySpec, ) @@ -372,7 +372,7 @@ def _get_update_properties(cls) -> list[PropertySpec]: ] -class AssetList(CogniteResourceList[Asset], IdTransformerMixin): +class AssetList(CogniteResourceList[Asset], IdAndExtIdTransformerMixin): _RESOURCE = Asset def __init__(self, resources: Collection[Any], cognite_client: CogniteClient | None = None) -> None: diff --git a/cognite/client/data_classes/data_sets.py b/cognite/client/data_classes/data_sets.py index ddc5c6080b..d75f329713 100644 --- a/cognite/client/data_classes/data_sets.py +++ b/cognite/client/data_classes/data_sets.py @@ -12,7 +12,7 @@ CogniteResource, CogniteResourceList, CogniteUpdate, - IdTransformerMixin, + IdAndExtIdTransformerMixin, PropertySpec, ) from cognite.client.data_classes.shared import TimestampRange @@ -171,5 +171,5 @@ def __init__(self, count: int | None = None, **kwargs: Any) -> None: count = CognitePropertyClassUtil.declare_property("count") -class DataSetList(CogniteResourceList[DataSet], IdTransformerMixin): +class DataSetList(CogniteResourceList[DataSet], IdAndExtIdTransformerMixin): _RESOURCE = DataSet diff --git a/cognite/client/data_classes/datapoints.py b/cognite/client/data_classes/datapoints.py index 7c1eb00aec..805ed36f75 100644 --- a/cognite/client/data_classes/datapoints.py +++ b/cognite/client/data_classes/datapoints.py @@ -657,65 +657,45 @@ def _repr_html_(self) -> str: return notebook_display_with_fallback(self, include_errors=is_synthetic_dps) -class DatapointsArrayList(CogniteResourceList[DatapointsArray]): - _RESOURCE = DatapointsArray - +class DuplicatedIdsMixin: def __init__(self, resources: Collection[Any], cognite_client: CogniteClient | None = None) -> None: super().__init__(resources, cognite_client) # Fix what happens for duplicated identifiers: ids = [dps.id for dps in self if dps.id is not None] xids = [dps.external_id for dps in self if dps.external_id is not None] - dupe_ids, id_dct = find_duplicates(ids), defaultdict(list) - dupe_xids, xid_dct = find_duplicates(xids), defaultdict(list) + idents_to_update = defaultdict(list) + dupe_ids, dupe_xids = find_duplicates(ids), find_duplicates(xids) for dps in self: if (id_ := dps.id) is not None and id_ in dupe_ids: - id_dct[id_].append(dps) + idents_to_update[("id", id_)].append(dps) if (xid := dps.external_id) is not None and xid in dupe_xids: - xid_dct[xid].append(dps) + idents_to_update[("external_id", xid)].append(dps) - self._id_to_item.update(id_dct) - self._external_id_to_item.update(xid_dct) + self._identifier_to_items.update(idents_to_update) - def concat_duplicate_ids(self) -> None: - """ - Concatenates all arrays with duplicated IDs. - Arrays with the same ids are stacked in chronological order. +class DatapointsArrayList(DuplicatedIdsMixin, CogniteResourceList[DatapointsArray]): + _RESOURCE = DatapointsArray - **Caveat** This method is not guaranteed to preserve the order of the list. - """ - # Rebuilt list instead of removing duplicated one at a time at the cost of O(n). + def concat_duplicate_ids(self) -> None: + """Concatenates all arrays with duplicated IDs in chronological order.""" + old_data = self.data[:] self.data.clear() - # This implementation takes advantage of the ordering of the duplicated in the __init__ method - has_external_ids = set() - for ext_id, items in self._external_id_to_item.items(): - if not isinstance(items, list): - self.data.append(items) - if items.id is not None: - has_external_ids.add(items.id) + for item in old_data: + if item.id is None: + raise ValueError("'concat_duplicate_ids' requires 'id' to be set on all elements") + if not isinstance(item, list): + self.data.append(item) continue - concatenated = DatapointsArray.create_from_arrays(*items) - self._external_id_to_item[ext_id] = concatenated - if concatenated.id is not None: - has_external_ids.add(concatenated.id) - self._id_to_item[concatenated.id] = concatenated - self.data.append(concatenated) - - if not (only_ids := set(self._id_to_item) - has_external_ids): - return - for id_, items in self._id_to_item.items(): - if id_ not in only_ids: - continue - if not isinstance(items, list): - self.data.append(items) - continue - concatenated = DatapointsArray.create_from_arrays(*items) - self._id_to_item[id_] = concatenated + concatenated = DatapointsArray.create_from_arrays(*item) self.data.append(concatenated) + self._identifier_to_items[("id", concatenated.id)] = concatenated.id + if concatenated.external_id is not None: + self._identifier_to_items[("external_id", concatenated.external_id)] = concatenated.external_id def get( # type: ignore [override] self, @@ -775,27 +755,9 @@ def dump(self, camel_case: bool = False, convert_timestamps: bool = False) -> li return [dps.dump(camel_case, convert_timestamps) for dps in self] -class DatapointsList(CogniteResourceList[Datapoints]): +class DatapointsList(DuplicatedIdsMixin, CogniteResourceList[Datapoints]): _RESOURCE = Datapoints - def __init__(self, resources: Collection[Any], cognite_client: CogniteClient | None = None) -> None: - super().__init__(resources, cognite_client) - - # Fix what happens for duplicated identifiers: - ids = [dps.id for dps in self if dps.id is not None] - xids = [dps.external_id for dps in self if dps.external_id is not None] - dupe_ids, id_dct = find_duplicates(ids), defaultdict(list) - dupe_xids, xid_dct = find_duplicates(xids), defaultdict(list) - - for dps in self: - if (id_ := dps.id) is not None and id_ in dupe_ids: - id_dct[id_].append(dps) - if (xid := dps.external_id) is not None and xid in dupe_xids: - xid_dct[xid].append(dps) - - self._id_to_item.update(id_dct) - self._external_id_to_item.update(xid_dct) - def get( # type: ignore [override] self, id: int | None = None, diff --git a/cognite/client/data_classes/documents.py b/cognite/client/data_classes/documents.py index 064f6de411..50bfd113fd 100644 --- a/cognite/client/data_classes/documents.py +++ b/cognite/client/data_classes/documents.py @@ -11,7 +11,7 @@ CogniteResourceList, CogniteSort, EnumProperty, - IdTransformerMixin, + IdAndExtIdTransformerMixin, NoCaseConversionPropertyList, ) from cognite.client.data_classes.aggregations import UniqueResult @@ -188,7 +188,7 @@ def dump(self, camel_case: bool = False) -> dict[str, Any]: return output -class DocumentList(CogniteResourceList[Document], IdTransformerMixin): +class DocumentList(CogniteResourceList[Document], IdAndExtIdTransformerMixin): _RESOURCE = Document diff --git a/cognite/client/data_classes/events.py b/cognite/client/data_classes/events.py index eb50b9dffa..1ddef7b1d1 100644 --- a/cognite/client/data_classes/events.py +++ b/cognite/client/data_classes/events.py @@ -16,7 +16,7 @@ CogniteSort, CogniteUpdate, EnumProperty, - IdTransformerMixin, + IdAndExtIdTransformerMixin, NoCaseConversionPropertyList, PropertySpec, ) @@ -251,7 +251,7 @@ def _get_update_properties(cls) -> list[PropertySpec]: ] -class EventList(CogniteResourceList[Event], IdTransformerMixin): +class EventList(CogniteResourceList[Event], IdAndExtIdTransformerMixin): _RESOURCE = Event diff --git a/cognite/client/data_classes/extractionpipelines.py b/cognite/client/data_classes/extractionpipelines.py index 5ea6dd6268..13d93b5c13 100644 --- a/cognite/client/data_classes/extractionpipelines.py +++ b/cognite/client/data_classes/extractionpipelines.py @@ -11,7 +11,7 @@ CogniteResource, CogniteResourceList, CogniteUpdate, - IdTransformerMixin, + IdAndExtIdTransformerMixin, PropertySpec, ) from cognite.client.data_classes.shared import TimestampRange @@ -212,7 +212,7 @@ def _get_update_properties(cls) -> list[PropertySpec]: ] -class ExtractionPipelineList(CogniteResourceList[ExtractionPipeline], IdTransformerMixin): +class ExtractionPipelineList(CogniteResourceList[ExtractionPipeline], IdAndExtIdTransformerMixin): _RESOURCE = ExtractionPipeline diff --git a/cognite/client/data_classes/files.py b/cognite/client/data_classes/files.py index dfd59f34b2..f60a33fcf7 100644 --- a/cognite/client/data_classes/files.py +++ b/cognite/client/data_classes/files.py @@ -12,7 +12,7 @@ CogniteResource, CogniteResourceList, CogniteUpdate, - IdTransformerMixin, + IdAndExtIdTransformerMixin, PropertySpec, ) from cognite.client.data_classes.labels import Label, LabelFilter @@ -297,5 +297,5 @@ def __init__(self, count: int | None = None, **kwargs: Any) -> None: count = CognitePropertyClassUtil.declare_property("count") -class FileMetadataList(CogniteResourceList[FileMetadata], IdTransformerMixin): +class FileMetadataList(CogniteResourceList[FileMetadata], IdAndExtIdTransformerMixin): _RESOURCE = FileMetadata diff --git a/cognite/client/data_classes/functions.py b/cognite/client/data_classes/functions.py index 810382f8b3..8e9af024f0 100644 --- a/cognite/client/data_classes/functions.py +++ b/cognite/client/data_classes/functions.py @@ -10,7 +10,7 @@ CogniteResource, CogniteResourceList, CogniteResponse, - IdTransformerMixin, + IdAndExtIdTransformerMixin, ) from cognite.client.data_classes.shared import TimestampRange from cognite.client.utils._time import ms_to_datetime @@ -263,7 +263,7 @@ class FunctionSchedulesList(CogniteResourceList[FunctionSchedule]): _RESOURCE = FunctionSchedule -class FunctionList(CogniteResourceList[Function], IdTransformerMixin): +class FunctionList(CogniteResourceList[Function], IdAndExtIdTransformerMixin): _RESOURCE = Function diff --git a/cognite/client/data_classes/raw.py b/cognite/client/data_classes/raw.py index eeb891c718..932e2741d2 100644 --- a/cognite/client/data_classes/raw.py +++ b/cognite/client/data_classes/raw.py @@ -3,7 +3,12 @@ from collections import OrderedDict from typing import TYPE_CHECKING, Any, cast, overload -from cognite.client.data_classes._base import CogniteResource, CogniteResourceList +from cognite.client.data_classes._base import ( + CogniteResource, + CogniteResourceList, + KeyTransformerMixin, + NameTransformerMixin, +) from cognite.client.utils._importing import local_import if TYPE_CHECKING: @@ -44,7 +49,7 @@ def to_pandas(self) -> pandas.DataFrame: # type: ignore[override] return pd.DataFrame([self.columns], [self.key]) -class RowList(CogniteResourceList[Row]): +class RowList(CogniteResourceList[Row], KeyTransformerMixin): _RESOURCE = Row def to_pandas(self) -> pandas.DataFrame: # type: ignore[override] @@ -106,7 +111,7 @@ def rows(self, key: str | None = None, limit: int | None = None) -> Row | RowLis return self._cognite_client.raw.rows.list(db_name=self._db_name, table_name=self.name, limit=limit) -class TableList(CogniteResourceList[Table]): +class TableList(CogniteResourceList[Table], NameTransformerMixin): _RESOURCE = Table @@ -143,5 +148,5 @@ def tables(self, limit: int | None = None) -> TableList: return self._cognite_client.raw.tables.list(db_name=self.name, limit=limit) -class DatabaseList(CogniteResourceList[Database]): +class DatabaseList(CogniteResourceList[Database], NameTransformerMixin): _RESOURCE = Database diff --git a/cognite/client/data_classes/sequences.py b/cognite/client/data_classes/sequences.py index 462ad1ab05..d60c87eefd 100644 --- a/cognite/client/data_classes/sequences.py +++ b/cognite/client/data_classes/sequences.py @@ -19,7 +19,7 @@ CogniteSort, CogniteUpdate, EnumProperty, - IdTransformerMixin, + IdAndExtIdTransformerMixin, NoCaseConversionPropertyList, PropertySpec, ) @@ -312,7 +312,7 @@ def __init__(self, count: int | None = None, **kwargs: Any) -> None: count = CognitePropertyClassUtil.declare_property("count") -class SequenceList(CogniteResourceList[Sequence], IdTransformerMixin): +class SequenceList(CogniteResourceList[Sequence], IdAndExtIdTransformerMixin): _RESOURCE = Sequence diff --git a/cognite/client/data_classes/time_series.py b/cognite/client/data_classes/time_series.py index 45804019e1..c71223c628 100644 --- a/cognite/client/data_classes/time_series.py +++ b/cognite/client/data_classes/time_series.py @@ -17,7 +17,7 @@ CogniteSort, CogniteUpdate, EnumProperty, - IdTransformerMixin, + IdAndExtIdTransformerMixin, NoCaseConversionPropertyList, PropertySpec, ) @@ -316,7 +316,7 @@ def __init__(self, count: int | None = None, **kwargs: Any) -> None: count = CognitePropertyClassUtil.declare_property("count") -class TimeSeriesList(CogniteResourceList[TimeSeries], IdTransformerMixin): +class TimeSeriesList(CogniteResourceList[TimeSeries], IdAndExtIdTransformerMixin): _RESOURCE = TimeSeries diff --git a/cognite/client/data_classes/user_profiles.py b/cognite/client/data_classes/user_profiles.py index fef6e17e55..311096453a 100644 --- a/cognite/client/data_classes/user_profiles.py +++ b/cognite/client/data_classes/user_profiles.py @@ -1,10 +1,9 @@ from __future__ import annotations import json -from typing import TYPE_CHECKING, Any, Sequence, cast +from typing import TYPE_CHECKING, Any, cast -from cognite.client.data_classes._base import CogniteResource, CogniteResourceList -from cognite.client.utils._text import to_camel_case +from cognite.client.data_classes._base import CogniteResource, CogniteResourceList, UserIdentTransformerMixin if TYPE_CHECKING: from cognite.client import CogniteClient @@ -49,30 +48,17 @@ def __init__( def _load(cls, resource: dict[str, Any] | str, cognite_client: CogniteClient | None = None) -> UserProfile: if isinstance(resource, str): resource = cast(dict[str, Any], json.loads(resource)) - to_load = { - "user_identifier": resource["userIdentifier"], - "last_updated_time": resource["lastUpdatedTime"], - } - for param in ["given_name", "surname", "email", "display_name", "job_title"]: - if (value := resource.get(to_camel_case(param))) is not None: - to_load[param] = value - return cls(**to_load, cognite_client=cognite_client) + return cls( + user_identifier=resource["userIdentifier"], + last_updated_time=resource["lastUpdatedTime"], + given_name=resource.get("givenName"), + surname=resource.get("surname"), + email=resource.get("email"), + display_name=resource.get("displayName"), + job_title=resource.get("jobTitle"), + cognite_client=cognite_client, + ) -class UserProfileList(CogniteResourceList[UserProfile]): +class UserProfileList(CogniteResourceList[UserProfile], UserIdentTransformerMixin): _RESOURCE = UserProfile - - def __init__(self, resources: Sequence[UserProfile], cognite_client: CogniteClient | None = None) -> None: - super().__init__(resources, cognite_client) - - del self._id_to_item, self._external_id_to_item - self._user_identifier_to_item = {item.user_identifier: item for item in self.data or []} - - def get(self, user_identifier: str) -> UserProfile | None: # type: ignore [override] - """Get an item from this list by user_identifier. - Args: - user_identifier (str): The user_identifier of the item to get. - Returns: - UserProfile | None: The requested item or None if not found. - """ - return self._user_identifier_to_item.get(user_identifier) diff --git a/cognite/client/data_classes/workflows.py b/cognite/client/data_classes/workflows.py index 6c602c96cf..faf6ca1f67 100644 --- a/cognite/client/data_classes/workflows.py +++ b/cognite/client/data_classes/workflows.py @@ -12,6 +12,7 @@ from cognite.client.data_classes._base import ( CogniteResource, CogniteResourceList, + ExtIdTransformerMixin, T_CogniteResource, ) from cognite.client.utils._text import convert_all_keys_to_snake_case, to_snake_case @@ -79,7 +80,7 @@ def __init__( self.created_time = created_time -class WorkflowList(CogniteResourceList[Workflow]): +class WorkflowList(CogniteResourceList[Workflow], ExtIdTransformerMixin): """This class represents a list of workflows.""" _RESOURCE = Workflow diff --git a/cognite/client/utils/_auxiliary.py b/cognite/client/utils/_auxiliary.py index 5d2e1ee367..0dda1eb931 100644 --- a/cognite/client/utils/_auxiliary.py +++ b/cognite/client/utils/_auxiliary.py @@ -3,6 +3,7 @@ import functools import math import numbers +import operator as op import platform import warnings from decimal import Decimal @@ -28,7 +29,7 @@ if TYPE_CHECKING: from cognite.client import CogniteClient - from cognite.client.data_classes._base import T_CogniteResource + from cognite.client.data_classes._base import T_CogniteResource, T_CogniteResourceList T = TypeVar("T") @@ -216,3 +217,12 @@ def rename_and_exclude_keys( aliases = aliases or {} exclude = exclude or set() return {aliases.get(k, k): v for k, v in dct.items() if k not in exclude} + + +def get_identifiers(res_lst: T_CogniteResourceList, attr: str, ignore_missing: bool) -> list: + if not res_lst: + return [] + identifiers = list(map(op.attrgetter(attr), res_lst)) + if not ignore_missing and None in identifiers: + raise ValueError(f"All {type(res_lst[0]).__name__} must have '{attr}', or pass 'ignore_missing=True'") + return identifiers diff --git a/tests/tests_integration/test_api/test_assets.py b/tests/tests_integration/test_api/test_assets.py index 1b82ef86a3..cbef93651b 100644 --- a/tests/tests_integration/test_api/test_assets.py +++ b/tests/tests_integration/test_api/test_assets.py @@ -452,7 +452,7 @@ def test_variable_number_of_root_assets(self, cognite_client, n_roots, root_test with create_hierarchy_with_cleanup(cognite_client, assets) as created: assert len(assets) == len(created) # Make sure `.get` has the exact same mapping keys: - assert set(AssetList(assets)._external_id_to_item) == set(created._external_id_to_item) + assert set(AssetList(assets).as_external_ids()) == set(created.as_external_ids()) @pytest.mark.parametrize( "n_id, n_xid, pass_hierarchy", @@ -470,14 +470,14 @@ def test_orphans__parent_linked_using_mixed_ids_xids( self, n_id, n_xid, pass_hierarchy, cognite_client, root_test_asset_subtree ): assets = generate_orphan_assets(n_id, n_xid, sample_from=root_test_asset_subtree) - expected = set(AssetList(assets)._external_id_to_item) + expected = set(AssetList(assets).as_external_ids()) if pass_hierarchy: assets = AssetHierarchy(assets, ignore_orphans=True) with create_hierarchy_with_cleanup(cognite_client, assets) as created: assert len(assets) == len(created) # Make sure `.get` has the exact same mapping keys: - assert expected == set(created._external_id_to_item) + assert expected == set(created.as_external_ids()) def test_orphans__blocked_if_passed_as_asset_hierarchy_instance(self, cognite_client, root_test_asset_subtree): assets = generate_orphan_assets(2, 2, sample_from=root_test_asset_subtree) @@ -607,7 +607,7 @@ def test_upsert_and_insert_in_same_request(self, cognite_client, upsert_mode, mo cognite_client, assets, upsert=True, upsert_mode=upsert_mode ) as patch_created: assert len(patch_created) == 4 - assert set(AssetList(assets)._external_id_to_item) == set(patch_created._external_id_to_item) + assert set(AssetList(assets).as_external_ids()) == set(patch_created.as_external_ids()) # 1+3 because 3 additional calls were made: # 1) Try create all (fail), 2) create non-duplicated (success), 3) update duplicated (success) assert 1 + 3 == cognite_client.assets._post.call_count diff --git a/tests/tests_integration/test_api/test_workflows.py b/tests/tests_integration/test_api/test_workflows.py index a922d64fe7..426c054d0b 100644 --- a/tests/tests_integration/test_api/test_workflows.py +++ b/tests/tests_integration/test_api/test_workflows.py @@ -35,7 +35,7 @@ def workflow_list(cognite_client: CogniteClient) -> WorkflowList: ) workflows = [workflow1, workflow2] listed = cognite_client.workflows.list() - existing = listed._external_id_to_item + existing = {name for name, value in listed._identifier_to_items} call_list = False for workflow in workflows: if workflow.external_id not in existing: