Skip to content

Commit de116c6

Browse files
erlendvollseterlendvollset
committed
Use instances.Properties class in records too
1 parent a55e0b7 commit de116c6

File tree

3 files changed

+67
-131
lines changed

3 files changed

+67
-131
lines changed

‎cognite/client/data_classes/data_modeling/ids.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,8 @@ def version(self) -> str | None: ...
183183
ConstraintIdentifier = tuple[ContainerId, str]
184184
IndexIdentifier = tuple[ContainerId, str]
185185
ViewIdentifier = ViewId | tuple[str, str] | tuple[str, str, str]
186+
SourceIdentifier = ContainerIdentifier | ViewIdentifier
187+
SourceId = ContainerId | ViewId
186188
DataModelIdentifier = DataModelId | tuple[str, str] | tuple[str, str, str]
187189
NodeIdentifier = NodeId | tuple[str, str, str]
188190
EdgeIdentifier = EdgeId | tuple[str, str, str]
@@ -198,6 +200,20 @@ def _load_space_identifier(ids: str | SequenceNotStr[str]) -> DataModelingIdenti
198200
)
199201

200202

203+
def _load_source_id(source_identifier: SourceIdentifier) -> SourceId:
204+
match source_identifier:
205+
case ContainerId() | ViewId():
206+
return source_identifier
207+
case (str(space), str(view_external_id), str(version)):
208+
return ViewId(space=space, external_id=view_external_id, version=version)
209+
case (str(space), str(container_external_id)):
210+
return ContainerId(space, container_external_id)
211+
raise ValueError(
212+
"Invalid source identifier format. Expected ContainerId or ViewId or tuple of "
213+
"strings((space, externalId) for container or (space, externalId, version) for views)."
214+
)
215+
216+
201217
def _load_identifier(
202218
ids: Id | Sequence[Id], id_type: Literal["container", "view", "data_model", "space", "node", "edge"]
203219
) -> DataModelingIdentifierSequence:

‎cognite/client/data_classes/data_modeling/instances.py‎

Lines changed: 49 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -59,8 +59,10 @@
5959
ContainerId,
6060
EdgeId,
6161
NodeId,
62+
SourceId,
63+
SourceIdentifier,
6264
ViewId,
63-
ViewIdentifier,
65+
_load_source_id,
6466
)
6567
from cognite.client.utils._auxiliary import exactly_one_is_not_none, find_duplicates, flatten_dict
6668
from cognite.client.utils._identifier import InstanceId
@@ -224,82 +226,93 @@ def _load(cls, resource: dict, cognite_client: CogniteClient | None = None) -> S
224226
_T = TypeVar("_T")
225227

226228

227-
class Properties(MutableMapping[ViewIdentifier, MutableMapping[PropertyIdentifier, PropertyValue]]):
228-
def __init__(self, properties: MutableMapping[ViewId, MutableMapping[PropertyIdentifier, PropertyValue]]) -> None:
229+
class Properties(MutableMapping[SourceIdentifier, MutableMapping[PropertyIdentifier, PropertyValue]]):
230+
def __init__(self, properties: MutableMapping[SourceId, MutableMapping[PropertyIdentifier, PropertyValue]]) -> None:
229231
self.data = properties
230232

231233
@classmethod
232234
def load(
233235
cls, data: MutableMapping[Space, MutableMapping[str, MutableMapping[PropertyIdentifier, PropertyValue]]]
234236
) -> Properties:
235-
props: MutableMapping[ViewId, MutableMapping[PropertyIdentifier, PropertyValue]] = {}
237+
props: MutableMapping[SourceId, MutableMapping[PropertyIdentifier, PropertyValue]] = {}
236238
for space, view_properties in data.items():
237-
for view_id_str, properties in view_properties.items():
238-
view_tuple = tuple(view_id_str.split("/", 1))
239-
if len(view_tuple) != 2:
240-
warnings.warn(
241-
f"Unknown type of view id: {view_id_str}, expected format <external_id>/<version>. Skipping...",
242-
stacklevel=2,
243-
)
244-
continue
245-
view_id = ViewId.load((space, *view_tuple))
246-
props[view_id] = properties
239+
for source_id_str, properties in view_properties.items():
240+
if "/" in source_id_str:
241+
view_tuple = tuple(source_id_str.split("/", 1))
242+
if len(view_tuple) != 2:
243+
warnings.warn(
244+
f"Unknown type of view id: {source_id_str}, expected format <external_id>/<version>. Skipping...",
245+
stacklevel=2,
246+
)
247+
continue
248+
source_id: SourceId = ViewId.load((space, *view_tuple))
249+
else:
250+
source_id = ContainerId.load((space, source_id_str))
251+
props[source_id] = properties
252+
247253
return cls(props)
248254

249255
def dump(self) -> dict[Space, dict[str, dict[PropertyIdentifier, PropertyValue]]]:
250256
props: dict[Space, dict[str, dict[PropertyIdentifier, PropertyValue]]] = defaultdict(dict)
251-
for view_id, properties in self.data.items():
252-
view_id_str = f"{view_id.external_id}/{view_id.version}"
253-
props[view_id.space][view_id_str] = cast(dict[PropertyIdentifier, PropertyValue], properties)
257+
for source_id, properties in self.data.items():
258+
if isinstance(source_id, ViewId):
259+
source_id_str = f"{source_id.external_id}/{source_id.version}"
260+
elif isinstance(source_id, ContainerId):
261+
source_id_str = source_id.external_id
262+
else:
263+
raise ValueError("SourceId must be either a ViewId or a ContainerId")
264+
props[source_id.space][source_id_str] = cast(dict[PropertyIdentifier, PropertyValue], properties)
254265
# Defaultdict is not yaml serializable
255266
return dict(props)
256267

257-
def items(self) -> ItemsView[ViewId, MutableMapping[PropertyIdentifier, PropertyValue]]:
268+
def items(self) -> ItemsView[SourceId, MutableMapping[PropertyIdentifier, PropertyValue]]:
258269
return self.data.items()
259270

260-
def keys(self) -> KeysView[ViewId]:
271+
def keys(self) -> KeysView[SourceId]:
261272
return self.data.keys()
262273

263274
def values(self) -> ValuesView[MutableMapping[PropertyIdentifier, PropertyValue]]:
264275
return self.data.values()
265276

266-
def __iter__(self) -> Iterator[ViewId]:
277+
def __iter__(self) -> Iterator[SourceId]:
267278
yield from self.keys()
268279

269-
def __getitem__(self, view: ViewIdentifier) -> MutableMapping[PropertyIdentifier, PropertyValue]:
270-
view_id = ViewId.load(view)
271-
return self.data.get(view_id, {})
280+
def __getitem__(self, item: SourceIdentifier) -> MutableMapping[PropertyIdentifier, PropertyValue]:
281+
source_id = _load_source_id(item)
282+
return self.data.get(source_id, {})
272283

273284
def __contains__(self, item: Any) -> bool:
274-
view_id = ViewId.load(item)
275-
return view_id in self.data
285+
source_id = _load_source_id(item)
286+
return source_id in self.data
276287

277288
@overload
278-
def get(self, view: ViewIdentifier) -> MutableMapping[PropertyIdentifier, PropertyValue] | None: ...
289+
def get(self, view: SourceIdentifier) -> MutableMapping[PropertyIdentifier, PropertyValue] | None: ...
279290

280291
@overload
281292
def get(
282-
self, view: ViewIdentifier, default: MutableMapping[PropertyIdentifier, PropertyValue] | _T
293+
self, view: SourceIdentifier, default: MutableMapping[PropertyIdentifier, PropertyValue] | _T
283294
) -> MutableMapping[PropertyIdentifier, PropertyValue] | _T: ...
284295

285296
def get(
286297
self,
287-
view: ViewIdentifier,
298+
view: SourceIdentifier,
288299
default: MutableMapping[PropertyIdentifier, PropertyValue] | None | _T | None = None,
289300
) -> MutableMapping[PropertyIdentifier, PropertyValue] | None | _T:
290-
view_id = ViewId.load(view)
291-
return self.data.get(view_id, default)
301+
source_id = _load_source_id(view)
302+
return self.data.get(source_id, default)
292303

293304
def __len__(self) -> int:
294305
return len(self.data)
295306

296-
def __delitem__(self, view: ViewIdentifier) -> None:
297-
view_id = ViewId.load(view)
298-
del self.data[view_id]
307+
def __delitem__(self, item: SourceIdentifier) -> None:
308+
source_id = _load_source_id(item)
309+
del self.data[source_id]
299310

300-
def __setitem__(self, view: ViewIdentifier, properties: MutableMapping[PropertyIdentifier, PropertyValue]) -> None:
301-
view_id = ViewId.load(view)
302-
self.data[view_id] = properties
311+
def __setitem__(
312+
self, item: SourceIdentifier, properties: MutableMapping[PropertyIdentifier, PropertyValue]
313+
) -> None:
314+
source_id = _load_source_id(item)
315+
self.data[source_id] = properties
303316

304317
def _repr_html_(self) -> str:
305318
pd = local_import("pandas")

‎cognite/client/data_classes/data_modeling/records.py‎

Lines changed: 2 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,13 @@
11
from __future__ import annotations
22

3-
from collections import defaultdict
4-
from collections.abc import ItemsView, Iterator, KeysView, MutableMapping, ValuesView
53
from dataclasses import dataclass
64
from datetime import datetime
7-
from typing import (
8-
TYPE_CHECKING,
9-
Any,
10-
TypeVar,
11-
cast,
12-
overload,
13-
)
5+
from typing import TYPE_CHECKING, Any
146

157
from typing_extensions import Self
168

179
from cognite.client.data_classes._base import CogniteObject, CogniteResource, CogniteResourceList
18-
from cognite.client.data_classes.data_modeling.ids import ContainerId, ContainerIdentifier
19-
from cognite.client.data_classes.data_modeling.instances import (
20-
PropertyIdentifier,
21-
PropertyValue,
22-
SourceData,
23-
)
10+
from cognite.client.data_classes.data_modeling.instances import Properties, SourceData
2411
from cognite.client.utils import datetime_to_ms
2512
from cognite.client.utils._text import convert_all_keys_to_camel_case
2613

@@ -75,86 +62,6 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
7562
return cls(id=id, sources=sources)
7663

7764

78-
_T = TypeVar("_T")
79-
80-
81-
# TODO: Use the Properties class from cognite.client.data_classes.data_modeling.instances, and make that support
82-
# both container ids and view ids. We need that for when we support containers for instances anyway.
83-
class Properties(MutableMapping[ContainerIdentifier, MutableMapping[PropertyIdentifier, PropertyValue]]):
84-
def __init__(
85-
self, properties: MutableMapping[ContainerId, MutableMapping[PropertyIdentifier, PropertyValue]]
86-
) -> None:
87-
self.data = properties
88-
89-
@classmethod
90-
def load(
91-
cls, data: MutableMapping[str, MutableMapping[str, MutableMapping[PropertyIdentifier, PropertyValue]]]
92-
) -> Properties:
93-
props: MutableMapping[ContainerId, MutableMapping[PropertyIdentifier, PropertyValue]] = {}
94-
for space, container_properties in data.items():
95-
for container_id_str, properties in container_properties.items():
96-
container_id = ContainerId.load((space, container_id_str))
97-
props[container_id] = properties
98-
return cls(props)
99-
100-
def dump(self) -> dict[str, dict[str, dict[PropertyIdentifier, PropertyValue]]]:
101-
props: dict[str, dict[str, dict[PropertyIdentifier, PropertyValue]]] = defaultdict(dict)
102-
for container_id, properties in self.data.items():
103-
extid = container_id.external_id
104-
props[container_id.space][extid] = cast(dict[PropertyIdentifier, PropertyValue], properties)
105-
# Defaultdict is not yaml serializable
106-
return dict(props)
107-
108-
def items(self) -> ItemsView[ContainerId, MutableMapping[PropertyIdentifier, PropertyValue]]:
109-
return self.data.items()
110-
111-
def keys(self) -> KeysView[ContainerId]:
112-
return self.data.keys()
113-
114-
def values(self) -> ValuesView[MutableMapping[PropertyIdentifier, PropertyValue]]:
115-
return self.data.values()
116-
117-
def __iter__(self) -> Iterator[ContainerId]:
118-
yield from self.keys()
119-
120-
def __getitem__(self, view: ContainerIdentifier) -> MutableMapping[PropertyIdentifier, PropertyValue]:
121-
view_id = ContainerId.load(view)
122-
return self.data.get(view_id, {})
123-
124-
def __contains__(self, item: Any) -> bool:
125-
view_id = ContainerId.load(item)
126-
return view_id in self.data
127-
128-
@overload
129-
def get(self, source: ContainerIdentifier) -> MutableMapping[PropertyIdentifier, PropertyValue] | None: ...
130-
131-
@overload
132-
def get(
133-
self, source: ContainerIdentifier, default: MutableMapping[PropertyIdentifier, PropertyValue] | _T
134-
) -> MutableMapping[PropertyIdentifier, PropertyValue] | _T: ...
135-
136-
def get(
137-
self,
138-
source: ContainerIdentifier,
139-
default: MutableMapping[PropertyIdentifier, PropertyValue] | None | _T | None = None,
140-
) -> MutableMapping[PropertyIdentifier, PropertyValue] | None | _T:
141-
source_id = ContainerId.load(source)
142-
return self.data.get(source_id, default)
143-
144-
def __len__(self) -> int:
145-
return len(self.data)
146-
147-
def __delitem__(self, source: ContainerIdentifier) -> None:
148-
source_id = ContainerId.load(source)
149-
del self.data[source_id]
150-
151-
def __setitem__(
152-
self, source: ContainerIdentifier, properties: MutableMapping[PropertyIdentifier, PropertyValue]
153-
) -> None:
154-
source_id = ContainerId.load(source)
155-
self.data[source_id] = properties
156-
157-
15865
@dataclass
15966
class Record(CogniteResource):
16067
id: RecordId

0 commit comments

Comments
 (0)