From 0492296b519f40d1817eb2e5972636c6a45b3186 Mon Sep 17 00:00:00 2001
From: Stefaan Lippens <stefaan.lippens@vito.be>
Date: Tue, 11 Jun 2024 11:36:39 +0200
Subject: [PATCH] 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
---
 openeogeotrellis/geopysparkdatacube.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/openeogeotrellis/geopysparkdatacube.py b/openeogeotrellis/geopysparkdatacube.py
index 4149e16c2..fb46d5a22 100644
--- 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