From a411109cbf909a316bac5ea2945c20b4a27c932d Mon Sep 17 00:00:00 2001
From: Jeroen Verstraelen <jeroen.verstraelen@vito.be>
Date: Tue, 21 Dec 2021 17:17:38 +0100
Subject: [PATCH] EP-3934: add chunk_polygon process

---
 openeo_driver/ProcessGraphDeserializer.py | 23 ++++++++++++++++++++++-
 openeo_driver/datacube.py                 |  3 +++
 openeo_driver/dry_run.py                  |  3 +++
 3 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/openeo_driver/ProcessGraphDeserializer.py b/openeo_driver/ProcessGraphDeserializer.py
index 67e75713..37dc1f52 100644
--- a/openeo_driver/ProcessGraphDeserializer.py
+++ b/openeo_driver/ProcessGraphDeserializer.py
@@ -14,7 +14,7 @@
 import openeo_processes
 import requests
 from dateutil.relativedelta import relativedelta
-from shapely.geometry import shape, mapping
+from shapely.geometry import shape, mapping, MultiPolygon
 
 import openeo.udf
 from openeo.capabilities import ComparableVersion
@@ -597,6 +597,27 @@ def reduce_dimension(args: dict, env: EvalEnv) -> DriverDataCube:
     return data_cube.reduce_dimension(reducer=reduce_pg, dimension=dimension, env=env)
 
 
+@process_registry_100.add_function
+def chunk_polygon(args: dict, env: EvalEnv) -> DriverDataCube:
+    reduce_pg = extract_deep(args, "process", "process_graph")
+    chunks = extract_arg(args, 'chunks')
+    data_cube = extract_arg(args, 'data')
+
+    if isinstance(chunks, DelayedVector):
+        polygons = list(chunks.geometries)
+        for p in polygons:
+            reason = "{m!s} is not a polygon.".format(m=p)
+            raise ProcessParameterInvalidException(parameter='chunks', process='chunk_polygon', reason=reason)
+        polygon = MultiPolygon(polygons)
+    else:
+        polygon = geojson_to_multipolygon(chunks)
+
+    if polygon.area == 0:
+        reason = "polygon {m!s} has an area of {a!r}".format(m=polygon, a=polygon.area)
+        raise ProcessParameterInvalidException(parameter='chunks', process='chunk_polygon', reason=reason)
+    return data_cube.chunk_polygon(reducer=reduce_pg, chunks=polygon, env=env)
+
+
 @process
 def add_dimension(args: dict, env: EvalEnv) -> DriverDataCube:
     data_cube = extract_arg(args, 'data')
diff --git a/openeo_driver/datacube.py b/openeo_driver/datacube.py
index bcfed74e..87bd9d55 100644
--- a/openeo_driver/datacube.py
+++ b/openeo_driver/datacube.py
@@ -48,6 +48,9 @@ def apply_tiles_spatiotemporal(self, process, context={}) -> 'DriverDataCube':
     def reduce_dimension(self, reducer, dimension: str, env: EvalEnv) -> 'DriverDataCube':
         self._not_implemented()
 
+    def chunk_polygon(self, reducer, chunks, env: EvalEnv, context={}) -> 'DriverDataCube':
+        self._not_implemented()
+
     def add_dimension(self, name: str, label, type: str = "other") -> 'DriverDataCube':
         self._not_implemented()
 
diff --git a/openeo_driver/dry_run.py b/openeo_driver/dry_run.py
index be5b131e..3dbd94bc 100644
--- a/openeo_driver/dry_run.py
+++ b/openeo_driver/dry_run.py
@@ -482,6 +482,9 @@ def reduce_dimension(self, reducer, dimension: str, env: EvalEnv) -> 'DryRunData
 
         return dc._process_metadata(self.metadata.reduce_dimension(dimension_name=dimension))
 
+    def chunk_polygon(self, reducer, chunks, env: EvalEnv, context={}) -> 'DryRunDataCube':
+        return self.filter_spatial(chunks)
+
     def add_dimension(self, name: str, label, type: str = "other") -> 'DryRunDataCube':
         return self._process_metadata(self.metadata.add_dimension(name=name, label=label, type=type))