Repository: Open-EO/openeo-python-client

Issue:
ProcessBuilder __getitem__ support
from https://github.com/openEOPlatform/documentation/issues/10

```python
def fit_function(x:ProcessBuilder, parameters):
    pi = math.pi
    a0 = array_element(parameters, 0)
    a1 = array_element(parameters, 1)
    a2 = array_element(parameters, 2)
    return a0 + a1*cos(2*pi/31557600*x) + a2*sin(2*pi/31557600*x) # 31557600 are the seconds in one year
```

It would be nice to allow writing `parameters[0]` instead of `array_element(parameters, 0)`

Discussion:
m-mohr: I like that this is supporting both indices and labels.

I think it would make sense to document that "integer" labels are not supported. (I have the same issue in the  JS client).
soxofaan: > I think it would make sense to document that "integer" labels are not supported. (

Indeed, a user however can still use integer labels by directly using `array_element`,  
this `__getitem__` trick is just syntactic sugar
m-mohr: Yes, I only meant this special case.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 831591f1..d475e078 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/openeo/internal/processes/generator.py b/openeo/internal/processes/generator.py
index 834cd900..a53dc321 100644
--- a/openeo/internal/processes/generator.py
+++ b/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 generate_process_py(processes: List[Process], output=sys.stdout, argv=None):
             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
diff --git a/openeo/processes.py b/openeo/processes.py
index 417984e6..3bbb21eb 100644
--- a/openeo/processes.py
+++ b/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 @@ class ProcessBuilder(ProcessBuilderBase):
     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


Test patch:
diff --git a/tests/rest/datacube/test_processbuilder.py b/tests/rest/datacube/test_processbuilder.py
index 9528a203..405aa97e 100644
--- a/tests/rest/datacube/test_processbuilder.py
+++ b/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
+        }
+    }