From 8f903974ca892fab9c7fbc6ec839cc74e5cb35bd Mon Sep 17 00:00:00 2001
From: Stefaan Lippens <stefaan.lippens@vito.be>
Date: Thu, 3 Aug 2023 10:20:12 +0200
Subject: [PATCH] fixup! Issue #114/#200 support promoting feature properties
 to cube values

---
 openeo_driver/datacube.py |  4 ++
 tests/test_vectorcube.py  | 84 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+)

diff --git a/openeo_driver/datacube.py b/openeo_driver/datacube.py
index d4cadf71..eaed398f 100644
--- a/openeo_driver/datacube.py
+++ b/openeo_driver/datacube.py
@@ -254,6 +254,7 @@ def from_geodataframe(
         data: gpd.GeoDataFrame,
         *,
         columns_for_cube: Union[List[str], str] = COLUMN_SELECTION_NUMERICAL,
+        # TODO: change default band name to "properties" (per `load_geojson`  spec introduced by https://github.com/Open-EO/openeo-processes/pull/427)
         dimension_name: str = DIM_BANDS,
     ) -> "DriverVectorCube":
         """
@@ -539,6 +540,9 @@ def geometry_count(self) -> int:
     def get_geometries(self) -> Sequence[shapely.geometry.base.BaseGeometry]:
         return self._geometries.geometry
 
+    def get_cube(self) -> Optional[xarray.DataArray]:
+        return self._cube
+
     def get_ids(self) -> Optional[Sequence]:
         return self._geometries.get("id")
 
diff --git a/tests/test_vectorcube.py b/tests/test_vectorcube.py
index def8322f..0647c87d 100644
--- a/tests/test_vectorcube.py
+++ b/tests/test_vectorcube.py
@@ -1,5 +1,6 @@
 import textwrap
 
+import geopandas
 import geopandas as gpd
 import numpy.testing
 import pyproj
@@ -157,6 +158,89 @@ def test_with_cube_to_geojson(self, gdf):
             }
         )
 
+    def test_from_geodataframe_default(self, gdf):
+        vc = DriverVectorCube.from_geodataframe(gdf)
+        assert vc.to_geojson() == DictSubSet(
+            {
+                "type": "FeatureCollection",
+                "features": [
+                    DictSubSet(
+                        {
+                            "type": "Feature",
+                            "properties": {"id": "first", "pop": 1234},
+                            "geometry": {
+                                "coordinates": (((1.0, 1.0), (3.0, 1.0), (2.0, 3.0), (1.0, 1.0)),),
+                                "type": "Polygon",
+                            },
+                        }
+                    ),
+                    DictSubSet(
+                        {
+                            "type": "Feature",
+                            "properties": {"id": "second", "pop": 5678},
+                            "geometry": {
+                                "coordinates": (((4.0, 2.0), (5.0, 4.0), (3.0, 4.0), (4.0, 2.0)),),
+                                "type": "Polygon",
+                            },
+                        }
+                    ),
+                ],
+            }
+        )
+        cube = vc.get_cube()
+        assert cube.dims == ("geometries", "bands")
+        assert cube.shape == (2, 1)
+        assert {k: list(v.values) for k, v in cube.coords.items()} == {"geometries": [0, 1], "bands": ["pop"]}
+
+    @pytest.mark.parametrize(
+        ["columns_for_cube", "expected"],
+        [
+            ("numerical", {"shape": (2, 1), "coords": {"geometries": [0, 1], "bands": ["pop"]}}),
+            ("all", {"shape": (2, 2), "coords": {"geometries": [0, 1], "bands": ["id", "pop"]}}),
+            ([], None),
+            (["id"], {"shape": (2, 1), "coords": {"geometries": [0, 1], "bands": ["id"]}}),
+            (["pop", "id"], {"shape": (2, 2), "coords": {"geometries": [0, 1], "bands": ["pop", "id"]}}),
+            # TODO: test specifying non-existent column (to be filled with no-data):
+            # (["pop", "nopenope"], {"shape": (2, 2), "coords": {"geometries": [0, 1], "bands": ["pop", "nopenope"]}}),
+        ],
+    )
+    def test_from_geodataframe_columns_for_cube(self, gdf, columns_for_cube, expected):
+        vc = DriverVectorCube.from_geodataframe(gdf, columns_for_cube=columns_for_cube)
+        assert vc.to_geojson() == DictSubSet(
+            {
+                "type": "FeatureCollection",
+                "features": [
+                    DictSubSet(
+                        {
+                            "type": "Feature",
+                            "properties": {"id": "first", "pop": 1234},
+                            "geometry": {
+                                "coordinates": (((1.0, 1.0), (3.0, 1.0), (2.0, 3.0), (1.0, 1.0)),),
+                                "type": "Polygon",
+                            },
+                        }
+                    ),
+                    DictSubSet(
+                        {
+                            "type": "Feature",
+                            "properties": {"id": "second", "pop": 5678},
+                            "geometry": {
+                                "coordinates": (((4.0, 2.0), (5.0, 4.0), (3.0, 4.0), (4.0, 2.0)),),
+                                "type": "Polygon",
+                            },
+                        }
+                    ),
+                ],
+            }
+        )
+        cube = vc.get_cube()
+        if expected is None:
+            assert cube is None
+        else:
+            assert cube.dims == ("geometries", "bands")
+            assert cube.shape == expected["shape"]
+            assert {k: list(v.values) for k, v in cube.coords.items()} == expected["coords"]
+
     @pytest.mark.parametrize(["geojson", "expected"], [
         (
                 {"type": "Polygon", "coordinates": [[(1, 1), (3, 1), (2, 3), (1, 1)]]},