Skip to content

Commit

Permalink
several improvements for typed instances (#1919)
Browse files Browse the repository at this point in the history
  • Loading branch information
@haakonvt
haakonvt authored Sep 16, 2024
1 parent dd79d3c commit 327467f
Show file tree
Hide file tree
Showing 6 changed files with 168 additions and 171 deletions.
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
---
repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.6.4
rev: v0.6.5
hooks:
- id: ruff
args:
Expand Down
65 changes: 34 additions & 31 deletions cognite/client/_api/data_modeling/instances.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,9 @@ def retrieve_edges(
Returns:
EdgeList[T_Edge] | T_Edge | Edge | None: The requested edges.
Retrieve nodes using a custom Edge class Flow
Retrieve edges using a custom typed class "Flow". Any property that you want to look up by a different attribute name,
e.g. you want `my_edge.flow_rate` to return the data for property `flowRate`, must use the PropertyOptions as shown below.
We strongly suggest you use snake_cased attribute names, as is done here:
>>> from cognite.client import CogniteClient
>>> from cognite.client.data_classes.data_modeling import EdgeId, TypedEdge, PropertyOptions, DirectRelationReference, ViewId
Expand All @@ -367,7 +369,7 @@ def retrieve_edges(
... flow_rate: float,
... start_node: DirectRelationReference,
... end_node: DirectRelationReference,
... deleted_time: Union[int, None] = None,
... deleted_time: int | None = None,
... ) -> None:
... super().__init__(
... space, external_id, version, type, last_updated_time, created_time, start_node, end_node, deleted_time
Expand All @@ -379,7 +381,9 @@ def retrieve_edges(
... return ViewId("sp_model_space", "flow", "1")
...
>>> client = CogniteClient()
>>> res = client.data_modeling.instances.retrieve_edges(NodeId("mySpace", "theFlow"), edge_cls=Flow)
>>> res = client.data_modeling.instances.retrieve_edges(
... EdgeId("mySpace", "theFlow"), edge_cls=Flow
... )
>>> isinstance(res, Flow)
"""
res = self._retrieve_typed(
Expand Down Expand Up @@ -437,7 +441,6 @@ def retrieve_nodes(
without providing a custom node class, but in that case, the retrieved nodes will be of the
built-in Node class.
Args:
nodes (NodeId | Sequence[NodeId] | tuple[str, str] | Sequence[tuple[str, str]]): Node id(s) to retrieve.
node_cls (type[T_Node]): The custom node class to use, the retrieved nodes will automatically be serialized to this class.
Expand All @@ -447,7 +450,10 @@ def retrieve_nodes(
Returns:
NodeList[T_Node] | T_Node | Node | None: The requested edges.
Retrieve nodes using a custom Node class Person
Retrieve nodes using a custom typed node class "Person". Any property that you want to look up by a different attribute name,
e.g. you want `my_node.birth_year` to return the data for property `birthYear`, must use the PropertyOptions as shown below.
We strongly suggest you use snake_cased attribute names, as is done here:
>>> from cognite.client import CogniteClient
>>> from cognite.client.data_classes.data_modeling import NodeId, TypedNode, PropertyOptions, DirectRelationReference, ViewId
Expand All @@ -462,9 +468,9 @@ def retrieve_nodes(
... last_updated_time: int,
... created_time: int,
... name: str,
... birth_year: Union[int, None] = None,
... type: Union[DirectRelationReference, None] = None,
... deleted_time: Union[int, None] = None,
... birth_year: int | None = None,
... type: DirectRelationReference | None = None,
... deleted_time: int | None = None,
... ):
... super().__init__(
... space=space,
Expand All @@ -483,7 +489,9 @@ def retrieve_nodes(
... return ViewId("myModelSpace", "Person", "1")
...
>>> client = CogniteClient()
>>> res = client.data_modeling.instances.retrieve_nodes(NodeId("myDataSpace", "myPerson"), node_cls=Person)
>>> res = client.data_modeling.instances.retrieve_nodes(
... NodeId("myDataSpace", "myPerson"), node_cls=Person
... )
>>> isinstance(res, Person)
"""
res = self._retrieve_typed(
Expand Down Expand Up @@ -856,8 +864,8 @@ def apply(
>>> from cognite.client import CogniteClient
>>> from cognite.client.data_classes.data_modeling import EdgeApply, NodeOrEdgeData, NodeApply
>>> client = CogniteClient()
>>> nodes = [NodeApply("mySpace", "myNodeId")]
>>> res = client.data_modeling.instances.apply(nodes)
>>> node = NodeApply("mySpace", "myNodeId")
>>> res = client.data_modeling.instances.apply(node)
Create two nodes with data with a one-to-many edge
Expand Down Expand Up @@ -922,24 +930,23 @@ def apply(
>>> my_date = datetime(2020, 3, 14, 15, 9, 26, 535000, tzinfo=timezone.utc)
>>> data_model_timestamp = datetime_to_ms_iso_timestamp(my_date) # "2020-03-14T15:09:26.535+00:00"
Create a typed node:
Create a typed node apply. Any property that you want to look up by a different attribute name, e.g. you want
`my_node.birth_year` to return the data for property `birthYear`, must use the PropertyOptions as shown below.
We strongly suggest you use snake_cased attribute names, as is done here:
>>> from cognite.client import CogniteClient
>>> from datetime import date
>>> from cognite.client.data_classes.data_modeling import TypedNodeApply, PropertyOptions
>>> class Person(TypedNodeApply):
... birth_date = PropertyOptions(identifier="birthDate")
>>> class PersonApply(TypedNodeApply):
... birth_year = PropertyOptions(identifier="birthYear")
...
... def __init__(self, space: str, external_id, name: str, birth_date: date):
... def __init__(self, space: str, external_id, name: str, birth_year: int):
... super().__init__(space, external_id, type=("sp_model_space", "Person"))
... self.name = name
... self.birth_date = birth_date
...
... self.birth_year = birth_year
... def get_source(self):
... return ViewId("sp_model_space", "Person", "v1")
...
>>> client = CogniteClient()
>>> person = Person("sp_date_space", "my_person", "John Doe", date(1980, 1, 1))
>>> person = PersonApply("sp_date_space", "my_person", "John Doe", 1980)
>>> res = client.data_modeling.instances.apply(nodes=person)
"""
other_parameters = {
Expand Down Expand Up @@ -1050,7 +1057,7 @@ def search(
Args:
view (ViewId): View to search in.
query (str | None): Query string that will be parsed and used for search.
instance_type (Literal["node", "edge"] | type[T_Node] | type[T_Edge]): Whether to search for nodes or edges.
instance_type (Literal["node", "edge"] | type[T_Node] | type[T_Edge]): Whether to search for nodes or edges. You can also pass a custom typed node (or edge class) inheriting from TypedNode (or TypedEdge). See apply, retrieve_nodes or retrieve_edges for an example.
properties (list[str] | None): Optional array of properties you want to search through. If you do not specify one or more properties, the service will search all text fields within the view.
target_units (list[TargetUnit] | None): Properties to convert to another unit. The API can only convert to another unit if a unit has been defined as part of the type on the underlying container being queried.
space (str | SequenceNotStr[str] | None): Restrict instance search to the given space (or list of spaces).
Expand Down Expand Up @@ -1095,7 +1102,7 @@ def search(
list_cls: type[NodeList[T_Node]] | type[EdgeList[T_Edge]] = NodeList[Node] # type: ignore[assignment]
resource_cls: type[Node] | type[Edge] = Node
elif instance_type == "edge":
list_cls = EdgeList # type: ignore[assignment]
list_cls = EdgeList[Edge] # type: ignore[assignment]
resource_cls = Edge
elif inspect.isclass(instance_type) and issubclass(instance_type, TypedNode):
list_cls = NodeList[T_Node]
Expand Down Expand Up @@ -1511,7 +1518,7 @@ def list(
"""`List instances <https://developer.cognite.com/api#tag/Instances/operation/advancedListInstance>`_
Args:
instance_type (Literal["node", "edge"] | type[T_Node] | type[T_Edge]): Whether to query for nodes or edges.
instance_type (Literal["node", "edge"] | type[T_Node] | type[T_Edge]): Whether to query for nodes or edges. You can also pass a custom typed node (or edge class) inheriting from TypedNode (or TypedEdge). See apply, retrieve_nodes or retrieve_edges for an example.
include_typing (bool): Whether to return property type information as part of the result.
sources (Source | Sequence[Source] | None): Views to retrieve properties from.
space (str | SequenceNotStr[str] | None): Only return instances in the given space (or list of spaces).
Expand Down Expand Up @@ -1569,15 +1576,11 @@ def list(
elif instance_type == "edge":
resource_cls, list_cls = _NodeOrEdgeResourceAdapter, EdgeList
elif inspect.isclass(instance_type) and issubclass(instance_type, TypedNode):
resource_cls, list_cls = (
_NodeOrEdgeResourceAdapter,
_TypedNodeOrEdgeListAdapter(instance_type), # type: ignore[assignment]
)
resource_cls = _NodeOrEdgeResourceAdapter
list_cls = _TypedNodeOrEdgeListAdapter(instance_type) # type: ignore[assignment]
elif inspect.isclass(instance_type) and issubclass(instance_type, TypedEdge):
resource_cls, list_cls = (
_NodeOrEdgeResourceAdapter(Node, instance_type), # type: ignore[assignment]
_TypedNodeOrEdgeListAdapter(instance_type), # type: ignore[assignment]
)
resource_cls = _NodeOrEdgeResourceAdapter(Node, instance_type) # type: ignore[assignment]
list_cls = _TypedNodeOrEdgeListAdapter(instance_type) # type: ignore[assignment]
else:
raise ValueError(f"Invalid instance type: {instance_type}")

Expand Down
11 changes: 5 additions & 6 deletions cognite/client/_api/data_modeling/views.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -268,10 +268,9 @@ def apply(self, view: ViewApply | Sequence[ViewApply]) -> View | ViewList:
... version="1",
... name="Movie",
... properties={
... "name": MappedPropertyApply(
... container=ContainerId("imdb", "Movie"),
... name="name",
... container_property_identifier="name",
... "title": MappedPropertyApply(
... container=ContainerId(space="imdb", external_id="Movie"),
... container_property_identifier="title",
... ),
... "actors": MultiEdgeConnectionApply(
... type=DirectRelationReference(
Expand All @@ -296,11 +295,11 @@ def apply(self, view: ViewApply | Sequence[ViewApply]) -> View | ViewList:
... ),
... "movies": MultiEdgeConnectionApply(
... type=DirectRelationReference(
... space="imdb", external_id="Movie.actors"
... space="imdb", external_id="Role.movies"
... ),
... source=ViewId("imdb", "Movie", "1"),
... name="movies",
... direction="inwards",
... direction="outwards",
... ),
... }
... )
Expand Down
7 changes: 2 additions & 5 deletions cognite/client/data_classes/data_modeling/data_types.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,12 +34,9 @@ def dump(self, camel_case: bool = True) -> dict[str, str | dict]:
@classmethod
def load(cls, data: dict | tuple[str, str] | DirectRelationReference) -> DirectRelationReference:
if isinstance(data, dict):
return cls(
space=data["space"],
external_id=data["externalId"],
)
return cls(space=data["space"], external_id=data["externalId"])
elif isinstance(data, tuple) and len(data) == 2:
return cls(data[0], data[1])
return cls(*data)
elif isinstance(data, cls):
return data
else:
Expand Down
Loading

0 comments on commit 327467f

Please sign in to comment.