Fixes #228 add ProcessBuilder array_element(data, n) through data[n] syntax

Author: soxofaan

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add `DataCube.dimension_labels()` (EP-4008)
 - Add `Connection.load_result()` (EP-4008)
 - Add proper support for child callbacks in `fit_curve` and `predict_curve` ([#229](https://github.com/Open-EO/openeo-python-client/issues/229))
+- `ProcessBuilder`: Add support for `array_element(data, n)` through `data[n]` syntax ([#228](https://github.com/Open-EO/openeo-python-client/issues/228))
 
 
 ### Changed

openeo/internal/processes/generator.py

@@ -107,6 +107,7 @@ def collect_processes(sources: List[Union[Path, str]]) -> List[Process]:
 
 def generate_process_py(processes: List[Process], output=sys.stdout, argv=None):
     oo_src = textwrap.dedent("""
+        import builtins
         from openeo.internal.processes.builder import ProcessBuilderBase, UNSET
         
         
@@ -142,6 +143,12 @@ def __neg__(self) -> 'ProcessBuilder':
             def __pow__(self, other) -> 'ProcessBuilder':
                 return self.power(other)
 
+            def __getitem__(self, key) -> 'ProcessBuilder':
+                if isinstance(key, builtins.int):
+                    return self.array_element(index=key)
+                else:
+                    return self.array_element(label=key)
+
     """)
     fun_src = textwrap.dedent("""
         # Public shortcut

openeo/processes.py

@@ -4,6 +4,7 @@
 # Used command line arguments:
 #    openeo/internal/processes/generator.py ../openeo-processes/ ../openeo-processes/proposals/ --output openeo/processes.py
 
+import builtins
 from openeo.internal.processes.builder import ProcessBuilderBase, UNSET
 
 
@@ -39,6 +40,12 @@ def __neg__(self) -> 'ProcessBuilder':
     def __pow__(self, other) -> 'ProcessBuilder':
         return self.power(other)
 
+    def __getitem__(self, key) -> 'ProcessBuilder':
+        if isinstance(key, builtins.int):
+            return self.array_element(index=key)
+        else:
+            return self.array_element(label=key)
+
     def absolute(self) -> 'ProcessBuilder':
         """
         Absolute value

tests/rest/datacube/test_processbuilder.py

@@ -275,10 +275,10 @@ def test_apply_dimension_bandmath_lambda(con100):
 
 
 def test_apply_dimension_time_to_bands(con100):
-    from openeo.processes import array_concat,quantiles,sd,mean
+    from openeo.processes import array_concat, quantiles, sd, mean
     im = con100.load_collection("S2")
     res = im.apply_dimension(
-        process=lambda d: array_concat(quantiles(d,[0.25,0.5,0.75]), [sd(d),mean(d)]),
+        process=lambda d: array_concat(quantiles(d, [0.25, 0.5, 0.75]), [sd(d), mean(d)]),
         dimension="t",
         target_dimension="bands"
     )
@@ -382,3 +382,81 @@ def test_merge_cubes_max_lambda(con100):
     im2 = con100.load_collection("MASK")
     res = im1.merge_cubes(other=im2, overlap_resolver=lambda data: data.max())
     assert res.graph == load_json_resource('data/1.0.0/merge_cubes_max.json')
+
+
+def test_getitem_array_element_index(con100):
+    im = con100.load_collection("S2")
+
+    def callback(data: ProcessBuilder):
+        return data[1] + data[2]
+
+    res = im.reduce_dimension(reducer=callback, dimension="bands")
+
+    assert res.flat_graph() == {
+        "loadcollection1": {
+            "process_id": "load_collection",
+            "arguments": {"id": "S2", "spatial_extent": None, "temporal_extent": None},
+        },
+        "reducedimension1": {
+            "process_id": "reduce_dimension",
+            "arguments": {
+                "data": {"from_node": "loadcollection1"},
+                "dimension": "bands",
+                "reducer": {"process_graph": {
+                    "arrayelement1": {
+                        "process_id": "array_element",
+                        "arguments": {"data": {"from_parameter": "data"}, "index": 1},
+                    },
+                    "arrayelement2": {
+                        "process_id": "array_element",
+                        "arguments": {"data": {"from_parameter": "data"}, "index": 2},
+                    },
+                    "add1": {
+                        "process_id": "add",
+                        "arguments": {"x": {"from_node": "arrayelement1"}, "y": {"from_node": "arrayelement2"}},
+                        "result": True
+                    },
+                }}
+            },
+            "result": True
+        }
+    }
+
+
+def test_getitem_array_element_label(con100):
+    im = con100.load_collection("S2")
+
+    def callback(data: ProcessBuilder):
+        return data["red"] + data["green"]
+
+    res = im.reduce_dimension(reducer=callback, dimension="bands")
+
+    assert res.flat_graph() == {
+        "loadcollection1": {
+            "process_id": "load_collection",
+            "arguments": {"id": "S2", "spatial_extent": None, "temporal_extent": None},
+        },
+        "reducedimension1": {
+            "process_id": "reduce_dimension",
+            "arguments": {
+                "data": {"from_node": "loadcollection1"},
+                "dimension": "bands",
+                "reducer": {"process_graph": {
+                    "arrayelement1": {
+                        "process_id": "array_element",
+                        "arguments": {"data": {"from_parameter": "data"}, "label": "red"},
+                    },
+                    "arrayelement2": {
+                        "process_id": "array_element",
+                        "arguments": {"data": {"from_parameter": "data"}, "label": "green"},
+                    },
+                    "add1": {
+                        "process_id": "add",
+                        "arguments": {"x": {"from_node": "arrayelement1"}, "y": {"from_node": "arrayelement2"}},
+                        "result": True
+                    },
+                }}
+            },
+            "result": True
+        }
+    }