Issue #732/#733 avoid extra save_result after datacube_from_flat_graph

quickfix for #732 and a first step for #733

Author: soxofaan

openeo/rest/connection.py

@@ -72,6 +72,7 @@
     ProcessListingResponse,
     ValidationResponse,
 )
+from openeo.rest.result import SaveResult
 from openeo.rest.service import Service
 from openeo.rest.udp import Parameter, RESTUserDefinedProcess
 from openeo.rest.userfile import UserFile
@@ -1037,9 +1038,12 @@ def datacube_from_process(self, process_id: str, namespace: Optional[str] = None
         graph = PGNode(process_id, namespace=namespace, arguments=kwargs)
         return DataCube(graph=graph, connection=self)
 
-    def datacube_from_flat_graph(self, flat_graph: dict, parameters: Optional[dict] = None) -> DataCube:
+    def datacube_from_flat_graph(
+        self, flat_graph: dict, parameters: Optional[dict] = None
+    ) -> Union[DataCube, SaveResult]:
         """
-        Construct a :py:class:`DataCube` from a flat dictionary representation of a process graph.
+        Construct a :py:class:`~openeo.rest.datacube.DataCube`/:py:class:`~openeo.rest.result.SaveResult`
+        from the operations encoded in a flat dictionary representation of an openEO process graph.
 
         .. seealso:: :ref:`datacube_from_json`, :py:meth:`~openeo.rest.connection.Connection.datacube_from_json`
 
@@ -1048,7 +1052,6 @@ def datacube_from_flat_graph(self, flat_graph: dict, parameters: Optional[dict]
             (and optionally parameter metadata under a "parameters" field).
         :param parameters: Optional dictionary mapping parameter names to parameter values
             to use for parameters occurring in the process graph (e.g. as used in user-defined processes)
-        :return: A :py:class:`DataCube` corresponding with the operations encoded in the process graph
         """
         parameters = parameters or {}
 
@@ -1062,18 +1065,26 @@ def datacube_from_flat_graph(self, flat_graph: dict, parameters: Optional[dict]
             flat_graph = flat_graph["process_graph"]
 
         pgnode = PGNode.from_flat_graph(flat_graph=flat_graph, parameters=parameters or {})
-        return DataCube(graph=pgnode, connection=self)
 
-    def datacube_from_json(self, src: Union[str, Path], parameters: Optional[dict] = None) -> DataCube:
+        # TODO #733: smarter or more dynamic detection of what kind of object to return?
+        # TODO #733: rename `datacube_from_flat_graph` to something more general as this is not only about DataCube's anymore?
+        if pgnode.process_id == "save_result":
+            return SaveResult(graph=pgnode, connection=self)
+        else:
+            return DataCube(graph=pgnode, connection=self)
+
+    def datacube_from_json(
+        self, src: Union[str, Path], parameters: Optional[dict] = None
+    ) -> Union[DataCube, SaveResult]:
         """
-        Construct a :py:class:`DataCube` from JSON resource containing (flat) process graph representation.
+        Construct a :py:class:`~openeo.rest.datacube.DataCube`/:py:class:`~openeo.rest.result.SaveResult`
+        from a JSON resource containing a (flat) openEO process graph representation.
 
         .. seealso:: :ref:`datacube_from_json`, :py:meth:`~openeo.rest.connection.Connection.datacube_from_flat_graph`
 
         :param src: raw JSON string, URL to JSON resource or path to local JSON file
         :param parameters: Optional dictionary mapping parameter names to parameter values
             to use for parameters occurring in the process graph (e.g. as used in user-defined processes)
-        :return: A :py:class:`DataCube` corresponding with the operations encoded in the process graph
         """
         return self.datacube_from_flat_graph(load_json_resource(src), parameters=parameters)
 

tests/rest/datacube/test_datacube100.py

@@ -3028,6 +3028,50 @@ def test_reduce_callback(self, con100):
         assert isinstance(cube, DataCube)
         assert cube.flat_graph() == flat_graph
 
+    @pytest.mark.parametrize(
+        ["flat_graph", "expected"],
+        [
+            (
+                {
+                    "lc": {"process_id": "load_collection", "arguments": {"id": "S2"}, "result": True},
+                },
+                {
+                    "loadcollection1": {"process_id": "load_collection", "arguments": {"id": "S2"}},
+                    "saveresult1": {
+                        "process_id": "save_result",
+                        "arguments": {"data": {"from_node": "loadcollection1"}, "format": "GTiff", "options": {}},
+                        "result": True,
+                    },
+                },
+            ),
+            (
+                {
+                    "lc": {"process_id": "load_collection", "arguments": {"id": "S2"}},
+                    "sr": {
+                        "process_id": "save_result",
+                        "arguments": {"data": {"from_node": "lc"}, "format": "netCDF"},
+                        "result": True,
+                    },
+                },
+                {
+                    "loadcollection1": {"process_id": "load_collection", "arguments": {"id": "S2"}},
+                    "saveresult1": {
+                        "process_id": "save_result",
+                        "arguments": {"data": {"from_node": "loadcollection1"}, "format": "netCDF"},
+                        "result": True,
+                    },
+                },
+            ),
+        ],
+    )
+    def test_datacube_from_flat_graph_save_result_handling(self, con100, dummy_backend, flat_graph, expected):
+        """
+        https://github.com/Open-EO/openeo-python-client/issues/732
+        """
+        cube = dummy_backend.connection.datacube_from_flat_graph(flat_graph)
+        cube.create_job()
+        assert dummy_backend.get_batch_pg() == expected
+
 
 def test_send_nan_json(con100, requests_mock):
     """https://github.com/Open-EO/openeo-python-client/issues/185"""