From 48d7f897a09465ab57b7e8e664c9cf6e7dedeb2f Mon Sep 17 00:00:00 2001
From: Xavier Olive <git@xoolive.org>
Date: Mon, 15 Apr 2024 15:27:42 +0200
Subject: [PATCH] fix set_geometry bug with geopandas

---
 src/traffic/core/airspace.py                  | 20 +++++++++++--------
 src/traffic/data/eurocontrol/ddr/airspaces.py |  1 +
 src/traffic/data/eurocontrol/ddr/freeroute.py |  1 +
 src/traffic/data/eurocontrol/eurofirs.py      |  1 +
 src/traffic/data/faa/_airspace_boundary.py    |  1 +
 5 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/traffic/core/airspace.py b/src/traffic/core/airspace.py
index e3ca8fdd..9cc9ecdb 100644
--- a/src/traffic/core/airspace.py
+++ b/src/traffic/core/airspace.py
@@ -124,14 +124,18 @@ def dissolve(self, keep_levels: bool = False) -> gpd.GeoDataFrame:
             columns += ["upper", "lower"]
         name_table = self.data[["designator", "name"]].drop_duplicates()
 
-        return gpd.GeoDataFrame(
-            self.data.groupby(columns)
-            .agg(dict(geometry=unary_union))
-            .reset_index()
-            .merge(name_table)
-        ).assign(  # centroid is only accessible on a GeoDataFrame...
-            longitude=lambda df: df.centroid.x,
-            latitude=lambda df: df.centroid.y,
+        return (
+            gpd.GeoDataFrame(
+                self.data.groupby(columns)
+                .agg(dict(geometry=unary_union))
+                .reset_index()
+                .merge(name_table)
+            )
+            .assign(  # centroid is only accessible on a GeoDataFrame...
+                longitude=lambda df: df.centroid.x,
+                latitude=lambda df: df.centroid.y,
+            )
+            .set_geometry("geometry")
         )
 
 
diff --git a/src/traffic/data/eurocontrol/ddr/airspaces.py b/src/traffic/data/eurocontrol/ddr/airspaces.py
index 1f2c9f8f..db5478ec 100644
--- a/src/traffic/data/eurocontrol/ddr/airspaces.py
+++ b/src/traffic/data/eurocontrol/ddr/airspaces.py
@@ -76,6 +76,7 @@ def init_cache(self) -> None:
             )
             .query("designator != component or geometry.notnull()")
             .assign(upper=lambda df: df.upper.replace(999, float("inf")))
+            .set_geometry("geometry")
         )
 
     def read_are(self, filename: Path) -> None:
diff --git a/src/traffic/data/eurocontrol/ddr/freeroute.py b/src/traffic/data/eurocontrol/ddr/freeroute.py
index e60535cd..cdf56ac5 100644
--- a/src/traffic/data/eurocontrol/ddr/freeroute.py
+++ b/src/traffic/data/eurocontrol/ddr/freeroute.py
@@ -55,6 +55,7 @@ def init_cache(self) -> None:
         self.read_sls(sls_file)
 
         data = gpd.GeoDataFrame.from_records(self.elements_list)
+        data = data.set_geometry("geometry")
         if "name" not in data.columns:
             data = data.assign(name="")
         self.fra = self.data = data.assign(type="FRA")
diff --git a/src/traffic/data/eurocontrol/eurofirs.py b/src/traffic/data/eurocontrol/eurofirs.py
index 761e4b3c..646bb326 100644
--- a/src/traffic/data/eurocontrol/eurofirs.py
+++ b/src/traffic/data/eurocontrol/eurofirs.py
@@ -35,6 +35,7 @@ def __init__(self, data: gpd.GeoDataFrame | None = None) -> None:
                     ),
                 )
                 .drop(columns=["UPPERUNIT", "LOWERUNIT", "EFFECTDATE", "ICAO"])
+                .set_geometry("geometry")
             )
 
 
diff --git a/src/traffic/data/faa/_airspace_boundary.py b/src/traffic/data/faa/_airspace_boundary.py
index 5b6c0a97..26d75dd9 100644
--- a/src/traffic/data/faa/_airspace_boundary.py
+++ b/src/traffic/data/faa/_airspace_boundary.py
@@ -40,6 +40,7 @@ def __init__(self, data: gpd.GeoDataFrame | None = None) -> None:
                     lower=lambda df: df.lower.replace(-9998, 0),
                     upper=lambda df: df.upper.replace(-9998, float("inf")),
                 )
+                .set_geometry("geometry")
             )
 
     def back(self) -> Dict[str, Airspace]: