fixup! Support UDF based spatial/temporal extents in load_collection/load_stac #831

Author: soxofaan

openeo/rest/connection.py

@@ -35,12 +35,7 @@
 import openeo
 from openeo.config import config_log, get_config_option
 from openeo.internal.documentation import openeo_process
-from openeo.internal.graph_building import (
-    FlatGraphableMixin,
-    PGNode,
-    _FromNodeMixin,
-    as_flat_graph,
-)
+from openeo.internal.graph_building import FlatGraphableMixin, PGNode, as_flat_graph
 from openeo.internal.jupyter import VisualDict, VisualList
 from openeo.internal.processes.builder import ProcessBuilderBase
 from openeo.internal.warnings import deprecated, legacy_alias
@@ -1191,8 +1186,8 @@ def load_collection(
         self,
         collection_id: Union[str, Parameter],
         spatial_extent: Union[dict, Parameter, shapely.geometry.base.BaseGeometry, str, Path, None] = None,
-        temporal_extent: Union[Sequence[InputDate], Parameter, str, _FromNodeMixin, None] = None,
-        bands: Union[Iterable[str], Parameter, str, _FromNodeMixin, None] = None,
+        temporal_extent: Union[Sequence[InputDate], Parameter, str, FlatGraphableMixin, None] = None,
+        bands: Union[Iterable[str], Parameter, str, FlatGraphableMixin, None] = None,
         properties: Union[
             Dict[str, Union[PGNode, Callable]], List[CollectionProperty], CollectionProperty, None
         ] = None,
@@ -1293,9 +1288,9 @@ def load_stac(
         self,
         url: str,
         spatial_extent: Union[
-            dict, Parameter, shapely.geometry.base.BaseGeometry, str, Path, _FromNodeMixin, None
+            dict, Parameter, shapely.geometry.base.BaseGeometry, str, Path, FlatGraphableMixin, None
         ] = None,
-        temporal_extent: Union[Sequence[InputDate], Parameter, str, _FromNodeMixin, None] = None,
+        temporal_extent: Union[Sequence[InputDate], Parameter, str, FlatGraphableMixin, None] = None,
         bands: Union[Iterable[str], Parameter, str, None] = None,
         properties: Union[
             Dict[str, Union[PGNode, Callable]], List[CollectionProperty], CollectionProperty, None

openeo/rest/datacube.py

@@ -39,7 +39,12 @@
 from openeo.api.process import Parameter, schema_supports
 from openeo.dates import get_temporal_extent
 from openeo.internal.documentation import openeo_process
-from openeo.internal.graph_building import PGNode, ReduceNode, _FromNodeMixin
+from openeo.internal.graph_building import (
+    FlatGraphableMixin,
+    PGNode,
+    ReduceNode,
+    _FromNodeMixin,
+)
 from openeo.internal.jupyter import in_jupyter_context
 from openeo.internal.processes.builder import (
     ProcessBuilderBase,
@@ -91,7 +96,7 @@
 
 
 # Type annotation aliases
-InputDate = Union[str, datetime.date, Parameter, PGNode, ProcessBuilderBase, _FromNodeMixin, None]
+InputDate = Union[str, datetime.date, Parameter, PGNode, ProcessBuilderBase, FlatGraphableMixin, None]
 
 
 class DataCube(_ProcessGraphAbstraction):
@@ -166,9 +171,9 @@ def load_collection(
         collection_id: Union[str, Parameter],
         connection: Optional[Connection] = None,
         spatial_extent: Union[
-            dict, Parameter, shapely.geometry.base.BaseGeometry, str, pathlib.Path, _FromNodeMixin, None
+            dict, Parameter, shapely.geometry.base.BaseGeometry, str, pathlib.Path, FlatGraphableMixin, None
         ] = None,
-        temporal_extent: Union[Sequence[InputDate], Parameter, str, _FromNodeMixin, None] = None,
+        temporal_extent: Union[Sequence[InputDate], Parameter, str, FlatGraphableMixin, None] = None,
         bands: Union[Iterable[str], Parameter, str, None] = None,
         fetch_metadata: bool = True,
         properties: Union[
@@ -482,22 +487,22 @@ def _get_temporal_extent(
         *args,
         start_date: InputDate = None,
         end_date: InputDate = None,
-        extent: Union[Sequence[InputDate], Parameter, str, _FromNodeMixin, None] = None,
-    ) -> Union[List[Union[str, Parameter, PGNode, _FromNodeMixin, None]], Parameter, _FromNodeMixin]:
+        extent: Union[Sequence[InputDate], Parameter, str, FlatGraphableMixin, None] = None,
+    ) -> Union[List[Union[str, Parameter, PGNode, FlatGraphableMixin, None]], Parameter, FlatGraphableMixin]:
         """Parameter aware temporal_extent normalizer"""
         # TODO: move this outside of DataCube class
         # TODO: return extent as tuple instead of list
-        if len(args) == 1 and isinstance(args[0], (Parameter, _FromNodeMixin)):
+        if len(args) == 1 and isinstance(args[0], (Parameter, FlatGraphableMixin)):
             assert start_date is None and end_date is None and extent is None
             return args[0]
-        elif len(args) == 0 and isinstance(extent, (Parameter, _FromNodeMixin)):
+        elif len(args) == 0 and isinstance(extent, (Parameter, FlatGraphableMixin)):
             assert start_date is None and end_date is None
             # TODO: warn about unexpected parameter schema
             return extent
         else:
             def convertor(d: Any) -> Any:
                 # TODO: can this be generalized through _FromNodeMixin?
-                if isinstance(d, Parameter) or isinstance(d, _FromNodeMixin):
+                if isinstance(d, Parameter) or isinstance(d, FlatGraphableMixin):
                     # TODO: warn about unexpected parameter schema
                     return d
                 elif isinstance(d, ProcessBuilderBase):
@@ -533,7 +538,7 @@ def filter_temporal(
         *args,
         start_date: InputDate = None,
         end_date: InputDate = None,
-        extent: Union[Sequence[InputDate], Parameter, str, _FromNodeMixin, None] = None,
+        extent: Union[Sequence[InputDate], Parameter, str, FlatGraphableMixin, None] = None,
     ) -> DataCube:
         """
         Limit the DataCube to a certain date range, which can be specified in several ways: