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}, + }, + ), + ], + } + )