quickfix for DriverVectorCube support in chunk_polygon

refs: Open-EO/openeo-python-driver#287, Open-EO/openeo-python-driver#288, Open-EO/openeo-python-driver#291 #801
Open-EO · Jun 11, 2024 · 0492296 · 0492296
1 parent c1b8676
commit 0492296
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/openeogeotrellis/geopysparkdatacube.py b/openeogeotrellis/geopysparkdatacube.py
@@ -584,7 +584,14 @@ def chunk_polygon(
 
         if isinstance(reducer, dict):
             reducer = GeoPySparkBackendImplementation.accept_process_graph(reducer)
-        chunks: List[Polygon] = chunks.geoms
+
+        if isinstance(chunks, Polygon):
+            chunks = [chunks]
+        elif isinstance(chunks, MultiPolygon):
+            chunks: List[Polygon] = chunks.geoms
+        else:
+            raise ValueError(f"Invalid type for `chunks`: {type(chunks)}")
+
         jvm = get_jvm()
 
         result_collection = None