Improve automatic adding of save_result

openeo/rest/_datacube.py

@@ -151,6 +151,55 @@ def _repr_html_(self):
         }
         return render_component("model-builder", data=process, parameters=parameters)
 
+    def execute(self) -> dict:
+        """Executes the process graph."""
+        return self._connection.execute(self.flat_graph())
+
+
+class _Cube(_ProcessGraphAbstraction):
+    """Common base class for raster and vector data cubes."""
+
+    _DEFAULT_OUTPUT_FORMAT = None
+
+    def _ensure_save_result(
+        self,
+        format: Optional[str] = None,
+        options: Optional[dict] = None,
+    ) -> "_Cube":
+        """
+        Make sure there is a (final) `save_result` node in the process graph.
+        If there is already one: check if it is consistent with the given format/options (if any)
+        and add a new one otherwise.
+
+        :param format: (optional) desired `save_result` file format
+        :param options: (optional) desired `save_result` file format parameters
+        :return:
+        """
+        # TODO: move to generic data cube parent class (not only for raster cubes, but also vector cubes)
+        result_node = self.result_node()
+        if result_node.process_id == "save_result":
+            # There is already a `save_result` node:
+            # check if it is consistent with given format/options (if any)
+            args = result_node.arguments
+            if format is not None and format.lower() != args["format"].lower():
+                raise ValueError(
+                    f"Existing `save_result` node with different format {args['format']!r} != {format!r}"
+                )
+            if options is not None and options != args["options"]:
+                raise ValueError(
+                    f"Existing `save_result` node with different options {args['options']!r} != {options!r}"
+                )
+            cube = self
+        else:
+            # No `save_result` node yet: automatically add it.
+            # TODO: define `save_result` and or `process` methods on `_Cube`?
+            cube = self.save_result(
+                format=format or self._DEFAULT_OUTPUT_FORMAT, options=options
+            )
+        return cube
+
+    # TODO: define `create_job`, `execute_batch`, ... here too
+
 
 class UDF:
     """

openeo/rest/datacube.py

@@ -30,7 +30,7 @@
 from openeo.metadata import CollectionMetadata, Band, BandDimension, TemporalDimension, SpatialDimension
 from openeo.processes import ProcessBuilder
 from openeo.rest import BandMathException, OperatorException, OpenEoClientException
-from openeo.rest._datacube import _ProcessGraphAbstraction, THIS, UDF
+from openeo.rest._datacube import _ProcessGraphAbstraction, THIS, UDF, _Cube
 from openeo.rest.job import BatchJob
 from openeo.rest.mlmodel import MlModel
 from openeo.rest.service import Service
@@ -49,8 +49,7 @@
 
 
 
-
-class DataCube(_ProcessGraphAbstraction):
+class DataCube(_Cube):
     """
     Class representing a openEO (raster) data cube.
 
@@ -59,7 +58,7 @@ class DataCube(_ProcessGraphAbstraction):
     """
 
     # TODO: set this based on back-end or user preference?
-    _DEFAULT_RASTER_FORMAT = "GTiff"
+    _DEFAULT_OUTPUT_FORMAT = "GTiff"
 
     def __init__(self, graph: PGNode, connection: 'openeo.Connection', metadata: CollectionMetadata = None):
         super().__init__(pgnode=graph, connection=connection)
@@ -1815,7 +1814,7 @@ def atmospheric_correction(
     @openeo_process
     def save_result(
         self,
-        format: str = _DEFAULT_RASTER_FORMAT,
+        format: str = _DEFAULT_OUTPUT_FORMAT,
         options: Optional[dict] = None,
     ) -> "DataCube":
         formats = set(self._connection.list_output_formats().keys())
@@ -1864,7 +1863,7 @@ def _ensure_save_result(
         else:
             # No `save_result` node yet: automatically add it.
             cube = self.save_result(
-                format=format or self._DEFAULT_RASTER_FORMAT, options=options
+                format=format or self._DEFAULT_OUTPUT_FORMAT, options=options
             )
         return cube
 
@@ -2009,8 +2008,9 @@ def save_user_defined_process(
         )
 
     def execute(self) -> dict:
-        """Executes the process graph of the imagery. """
-        return self._connection.execute(self.flat_graph())
+        """Executes the process graph."""
+        # TODO: still necessary to do this explicitly here?
+        return super().execute()
 
     @staticmethod
     @deprecated(reason="Use :py:func:`openeo.udf.run_code.execute_local_udf` instead", version="0.7.0")

openeo/rest/mlmodel.py

@@ -24,6 +24,8 @@ class MlModel(_ProcessGraphAbstraction):
 
     .. versionadded:: 0.10.0
     """
+
+    # TODO
     def __init__(self, graph: PGNode, connection: 'Connection'):
         super().__init__(pgnode=graph, connection=connection)
 

openeo/rest/vectorcube.py

@@ -6,7 +6,7 @@
 from openeo.internal.graph_building import PGNode
 from openeo.internal.warnings import legacy_alias
 from openeo.metadata import CollectionMetadata
-from openeo.rest._datacube import _ProcessGraphAbstraction, UDF
+from openeo.rest._datacube import _ProcessGraphAbstraction, UDF, _Cube
 from openeo.rest.mlmodel import MlModel
 from openeo.rest.job import BatchJob
 from openeo.util import dict_no_none
@@ -16,7 +16,7 @@
     from openeo import Connection
 
 
-class VectorCube(_ProcessGraphAbstraction):
+class VectorCube(_Cube):
     """
     A Vector Cube, or 'Vector Collection' is a data structure containing 'Features':
     https://www.w3.org/TR/sdw-bp/#dfn-feature
@@ -90,6 +90,10 @@ def run_udf(
 
     @openeo_process
     def save_result(self, format: str = "GeoJson", options: dict = None):
+        # TODO?
+        # TODO: check format against supported formats
+        # TODO: should not return a VectorCube again, but bool wrapper
+        # TODO: should save_result also work on non-cube data types, e.g. arrays, scalars?
         return self.process(
             process_id="save_result",
             arguments={
@@ -105,6 +109,7 @@ def execute(self) -> dict:
 
     def download(self, outputfile: str, format: str = "GeoJSON", options: dict = None):
         # TODO: only add save_result, when not already present (see DataCube.download)
+        # TODO
         cube = self.save_result(format=format, options=options)
         return self._connection.download(cube.flat_graph(), outputfile)