Added chain to and direction argument to expression result sets (#1528)

CHANGELOG.md

@@ -17,6 +17,10 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## [7.5.0] - 2023-11-30
+### Added
+- `chain_to` to `NodeResultSetExpression` and `NodeResultSetExpression`, and `direction` to `NodeResultSetExpression`.
+
 ## [7.4.2] - 2023-11-28
 ### Improved
 - Quality of life improvement to `client.extraction_pipelines.runs.list` method. The `statuses` parameter now accepts

cognite/client/_version.py

@@ -1,4 +1,4 @@
 from __future__ import annotations
 
-__version__ = "7.4.2"
+__version__ = "7.5.0"
 __api_subversion__ = "V20220125"

cognite/client/data_classes/data_modeling/query.py

@@ -140,11 +140,21 @@ def __eq__(self, other: Any) -> bool:
 
 
 class ResultSetExpression(CogniteObject, ABC):
-    def __init__(self, from_: str | None, filter: Filter | None, limit: int | None, sort: list[InstanceSort] | None):
+    def __init__(
+        self,
+        from_: str | None,
+        filter: Filter | None,
+        limit: int | None,
+        sort: list[InstanceSort] | None,
+        direction: Literal["outwards", "inwards"] = "outwards",
+        chain_to: Literal["destination", "source"] = "destination",
+    ):
         self.from_ = from_
         self.filter = filter
         self.limit = limit
         self.sort = sort
+        self.direction = direction
+        self.chain_to = chain_to
 
     @abstractmethod
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
@@ -162,6 +172,8 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
             node = {
                 "from_": query_node.get("from"),
                 "filter": Filter.load(query_node["filter"]) if "filter" in query_node else None,
+                "chain_to": query_node.get("chainTo"),
+                "direction": query_node.get("direction"),
             }
             if (through := query_node.get("through")) is not None:
                 node["through"] = [
@@ -182,6 +194,7 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
                 if "terminationFilter" in query_edge
                 else None,
                 "limit_each": query_edge.get("limitEach"),
+                "chain_to": query_edge.get("chainTo"),
             }
             post_sort = [InstanceSort.load(sort) for sort in resource["postSort"]] if "postSort" in resource else []
             return cast(
@@ -195,15 +208,36 @@ def __eq__(self, other: Any) -> bool:
 
 
 class NodeResultSetExpression(ResultSetExpression):
+    """Describes how to query for nodes in the data model.
+
+    Args:
+        from_ (str | None): Chain your result-expression based on this view.
+        filter (Filter | None): Filter the result set based on this filter.
+        sort (list[InstanceSort] | None): Sort the result set based on this list of sort criteria.
+        limit (int | None): Limit the result set to this number of instances.
+        through (list[str] | tuple[str, str, str] | None): Chain your result-expression through this view.
+            The tuple must be on the form (space, view/version, property).
+        direction (Literal["outwards", "inwards"]): The direction to use when traversing direct relations.
+            Only applicable when through is specified.
+        chain_to (Literal["destination", "source"]): Control which side of the edge to chain to.
+            The chain_to option is only applicable if the result rexpression referenced in `from`
+            contains edges. `source` will chain to start if you're following edges outwards i.e `direction=outwards`. If you're
+            following edges inwards i.e `direction=inwards`, it will chain to end. `destination` (default) will chain to
+            end if you're following edges outwards i.e `direction=outwards`. If you're following edges
+            inwards i.e, `direction=inwards`, it will chain to start.
+    """
+
     def __init__(
         self,
         from_: str | None = None,
         filter: Filter | None = None,
         sort: list[InstanceSort] | None = None,
         limit: int | None = None,
         through: list[str] | tuple[str, str, str] | None = None,
-    ):
-        super().__init__(from_=from_, filter=filter, limit=limit, sort=sort)
+        direction: Literal["outwards", "inwards"] = "outwards",
+        chain_to: Literal["destination", "source"] = "destination",
+    ) -> None:
+        super().__init__(from_=from_, filter=filter, limit=limit, sort=sort, direction=direction, chain_to=chain_to)
         self.through = through
 
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
@@ -225,6 +259,10 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
                 },
                 "identifier": self.through[2],
             }
+        if self.chain_to:
+            nodes["chainTo" if camel_case else "chain_to"] = self.chain_to
+        if self.direction:
+            nodes["direction"] = self.direction
 
         if self.sort:
             output["sort"] = [s.dump(camel_case=camel_case) for s in self.sort]
@@ -235,6 +273,30 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
 
 
 class EdgeResultSetExpression(ResultSetExpression):
+    """Describes how to query for edges in the data model.
+
+    Args:
+        from_ (str | None): Chain your result expression from this edge.
+        max_distance (int | None): The largest - max - number of levels to traverse.
+        direction (Literal["outwards", "inwards"]): The direction to use when traversing.
+        filter (Filter | None): Filter the result set based on this filter.
+        node_filter (Filter | None): Filter the result set based on this filter.
+        termination_filter (Filter | None): Filter the result set based on this filter.
+        limit_each (int | None): Limit the number of returned edges for each of the source nodes in the result set.
+            The indicated uniform limit applies to the result set from the referenced from.
+            limitEach only has meaning when you also specify maxDistance=1 and from.
+        sort (list[InstanceSort] | None): Sort the result set based on this list of sort criteria.
+        post_sort (list[InstanceSort] | None): Sort the result set based on this list of sort criteria.
+        limit (int | None): Limit the result set to this number of instances.
+        chain_to (Literal["destination", "source"]): Control which side of the edge to chain to.
+            The chain_to option is only applicable if the result rexpression referenced in `from`
+            contains edges. `source` will chain to start if you're following edges outwards i.e `direction=outwards`. If you're
+            following edges inwards i.e `direction=inwards`, it will chain to end. `destination` (default) will chain to
+            end if you're following edges outwards i.e `direction=outwards`. If you're following edges
+            inwards i.e, `direction=inwards`, it will chain to start.
+
+    """
+
     def __init__(
         self,
         from_: str | None = None,
@@ -247,10 +309,10 @@ def __init__(
         sort: list[InstanceSort] | None = None,
         post_sort: list[InstanceSort] | None = None,
         limit: int | None = None,
-    ):
-        super().__init__(from_=from_, filter=filter, limit=limit, sort=sort)
+        chain_to: Literal["destination", "source"] = "destination",
+    ) -> None:
+        super().__init__(from_=from_, filter=filter, limit=limit, sort=sort, direction=direction, chain_to=chain_to)
         self.max_distance = max_distance
-        self.direction = direction
         self.node_filter = node_filter
         self.termination_filter = termination_filter
         self.limit_each = limit_each
@@ -273,6 +335,8 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
             edges["terminationFilter" if camel_case else "termination_filter"] = self.termination_filter.dump()
         if self.limit_each:
             edges["limitEach" if camel_case else "limit_each"] = self.limit_each
+        if self.chain_to:
+            edges["chainTo" if camel_case else "chain_to"] = self.chain_to
 
         if self.sort:
             output["sort"] = [s.dump(camel_case=camel_case) for s in self.sort]

pyproject.toml

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

tests/tests_unit/test_data_classes/test_data_models/test_queries.py

@@ -17,6 +17,8 @@ def result_set_expression_load_and_dump_equals_data() -> Iterator[ParameterSet]:
                 "view": {"type": "view", "space": "some", "externalId": "extid", "version": "v1"},
                 "identifier": "bla",
             },
+            "chainTo": "destination",
+            "direction": "outwards",
         },
         "limit": 1,
     }
@@ -30,7 +32,9 @@ def result_set_expression_load_and_dump_equals_data() -> Iterator[ParameterSet]:
 
     raw = {
         "nodes": {
-            "filter": {"range": {"lt": 2000, "property": ["IntegrationTestsImmutable", "Movie/2", "releaseYear"]}}
+            "filter": {"range": {"lt": 2000, "property": ["IntegrationTestsImmutable", "Movie/2", "releaseYear"]}},
+            "chainTo": "destination",
+            "direction": "outwards",
         }
     }
     loaded_node = q.NodeResultSetExpression(
@@ -45,6 +49,7 @@ def result_set_expression_load_and_dump_equals_data() -> Iterator[ParameterSet]:
             "filter": {
                 "equals": {"property": ["edge", "type"], "value": {"space": "MovieSpace", "externalId": "Movie.actors"}}
             },
+            "chainTo": "destination",
         }
     }
     loaded_edge = q.EdgeResultSetExpression(
@@ -124,6 +129,8 @@ def query_load_yaml_data() -> Iterator[ParameterSet]:
                 equals:
                     property: ["node", "externalId"]
                     value: {"parameter": "airplaneExternalId"}
+            chainTo: destination
+            direction: outwards
         limit: 1
     lands_in_airports:
         edges:
@@ -134,8 +141,11 @@ def query_load_yaml_data() -> Iterator[ParameterSet]:
                 equals:
                     property: ["edge", "type"]
                     value: ["aviation", "lands-in"]
+            chainTo: destination
     airports:
         nodes:
+            chainTo: destination
+            direction: outwards
             from: lands_in_airports
 parameters:
     airplaneExternalId: myFavouriteAirplane
@@ -171,6 +181,8 @@ def query_load_yaml_data() -> Iterator[ParameterSet]:
           - Movie/2
           - releaseYear
           value: 1994
+      chainTo: destination
+      direction: outwards
 select:
   movies:
     sources:

tests/utils.py

@@ -308,12 +308,12 @@ def create_instance(self, resource_cls: type[T_Object]) -> T_Object:
             # DataPointSubscriptionCreate requires either timeseries_ids or filter
             keyword_arguments.pop("filter", None)
         if resource_cls is Query:
-            # keys in with must match keys in select
-            selects = list(keyword_arguments["select"].values())
-            new_selects = {}
-            for i, key in enumerate(keyword_arguments["with_"]):
-                new_selects[key] = selects[i]
-            keyword_arguments["select"] = new_selects
+            # The fake generator makes all dicts from 1-3 values, we need to make sure that the query is valid
+            # by making sure that the list of equal length, so we make both to length 1.
+            with_key, with_value = next(iter(keyword_arguments["with_"].items()))
+            select_value = next(iter(keyword_arguments["select"].values()))
+            keyword_arguments["with_"] = {with_key: with_value}
+            keyword_arguments["select"] = {with_key: select_value}
         elif resource_cls is Relationship:
             # Relationship must set the source and target type consistently with the source and target
             keyword_arguments["source_type"] = type(keyword_arguments["source"]).__name__