From 09800babb0546e49c1e9cfbbd362a8a15a42df4c Mon Sep 17 00:00:00 2001
From: Stefaan Lippens <stefaan.lippens@vito.be>
Date: Tue, 6 Jun 2023 11:02:29 +0200
Subject: [PATCH] Issue #197 Initial support for run_udf on DriverVectorCube

---
 openeo_driver/datacube.py | 23 +++++++++++++++-
 tests/test_vectorcube.py  | 55 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/openeo_driver/datacube.py b/openeo_driver/datacube.py
index f84b89ad..32ec5407 100644
--- a/openeo_driver/datacube.py
+++ b/openeo_driver/datacube.py
@@ -16,6 +16,8 @@
 from pyproj import CRS
 import requests
 
+import openeo.udf
+import openeo.udf.run_code
 from openeo.metadata import CollectionMetadata
 from openeo.util import ensure_dir
 from openeo_driver.datastructs import SarBackscatterArgs, ResolutionMergeArgs, StacAsset
@@ -200,7 +202,7 @@ def __init__(self, message="Unspecified VectorCube error"):
         super(VectorCubeError, self).__init__(message=message)
 
 
-class DriverVectorCube:
+class DriverVectorCube(SupportsRunUdf):
     """
     Base class for driver-side 'vector cubes'
 
@@ -500,6 +502,25 @@ def buffer_points(self, distance: float = 10) -> "DriverVectorCube":
             ]
         )
 
+    def supports_udf(self, udf: str, runtime: str = "Python") -> bool:
+        udf_globals = openeo.udf.run_code.load_module_from_string(code=udf)
+        return any(
+            name in udf_globals
+            for name in [
+                "udf_apply_geojson_feature_collection",
+            ]
+        )
+
+    def run_udf(self, *, udf: str, runtime: str = "Python", context: Optional[dict] = None, env: EvalEnv):
+        udf_globals = openeo.udf.run_code.load_module_from_string(code=udf)
+
+        if "udf_apply_geojson_feature_collection" in udf_globals:
+            callback = udf_globals["udf_apply_geojson_feature_collection"]
+            result = callback(self.to_geojson())
+            return DriverVectorCube.from_geojson(geojson=result)
+        else:
+            raise openeo.udf.OpenEoUdfException("No UDF found")
+
 
 class DriverMlModel:
     """Base class for driver-side 'ml-model' data structures"""
diff --git a/tests/test_vectorcube.py b/tests/test_vectorcube.py
index 7c61a0d9..23d3a445 100644
--- a/tests/test_vectorcube.py
+++ b/tests/test_vectorcube.py
@@ -1,3 +1,5 @@
+import textwrap
+
 import geopandas as gpd
 import numpy.testing
 import pyproj
@@ -446,3 +448,56 @@ def test_buffer_points(self):
                 ],
             }
         )
+
+    def test_run_udf_geojson(self, gdf):
+        vc = DriverVectorCube(gdf)
+        udf = textwrap.dedent(
+            """
+            def udf_apply_geojson_feature_collection(feature_collection: dict) -> dict:
+                features = []
+                for feature in feature_collection["features"]:
+                    # transpose and rescale vertices
+                    vertices = feature["geometry"]["coordinates"][0]
+                    vertices = [(y, 2 * x) for (x, y) in vertices]
+                    feature["geometry"]["coordinates"] = [vertices]
+
+                    # Manipulate properties
+                    feature["properties"]["hello"] = feature["properties"]["id"]
+                    feature["properties"]["vertices"] = len(vertices) - 1
+                    del feature["properties"]["pop"]
+
+                    features.append(feature)
+                return {"type": "FeatureCollection", "features": features}
+        """
+        )
+        vc2 = vc.run_udf(udf=udf)
+        assert isinstance(vc2, DriverVectorCube)
+        assert vc2.to_geojson() == DictSubSet(
+            {
+                "type": "FeatureCollection",
+                "features": [
+                    DictSubSet(
+                        {
+                            "type": "Feature",
+                            "id": "0",
+                            "geometry": {
+                                "type": "Polygon",
+                                "coordinates": (((1, 2), (1, 6), (3, 4), (1, 2)),),
+                            },
+                            "properties": {"hello": "first", "id": "first", "vertices": 3},
+                        }
+                    ),
+                    DictSubSet(
+                        {
+                            "type": "Feature",
+                            "id": "1",
+                            "geometry": {
+                                "type": "Polygon",
+                                "coordinates": (((2, 8), (4, 10), (4, 6), (2, 8)),),
+                            },
+                            "properties": {"hello": "second", "id": "second", "vertices": 3},
+                        },
+                    ),
+                ],
+            }
+        )