From 3d3faf7997f719e2eb94323f7b044893681b7394 Mon Sep 17 00:00:00 2001
From: Krasen Samardzhiev <krasensam@gmail.com>
Date: Mon, 17 Jun 2024 12:47:40 +0200
Subject: [PATCH] BUG: fix describe_ function when count is not present (#608)

---
 momepy/functional/_diversity.py           |  6 ++++--
 momepy/functional/tests/test_diversity.py | 25 +++++++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/momepy/functional/_diversity.py b/momepy/functional/_diversity.py
index e5240e70..6efb020b 100644
--- a/momepy/functional/_diversity.py
+++ b/momepy/functional/_diversity.py
@@ -164,7 +164,8 @@ def describe_agg(
     result.loc[stats.index.values] = stats.values
     result.columns = stats.columns
     # fill only counts with zeros, other stats are NA
-    result.loc[:, "count"] = result.loc[:, "count"].fillna(0)
+    if "count" in result.columns:
+        result.loc[:, "count"] = result.loc[:, "count"].fillna(0)
     result.index.names = result_index.names
 
     return result
@@ -263,7 +264,8 @@ def describe_reached_agg(
     result.loc[stats.index.values] = stats.values
     result.columns = stats.columns
     # fill only counts with zeros, other stats are NA
-    result.loc[:, "count"] = result.loc[:, "count"].fillna(0)
+    if "count" in result.columns:
+        result.loc[:, "count"] = result.loc[:, "count"].fillna(0)
     result.index.name = None
 
     return result
diff --git a/momepy/functional/tests/test_diversity.py b/momepy/functional/tests/test_diversity.py
index a0055cbc..2d28000d 100644
--- a/momepy/functional/tests/test_diversity.py
+++ b/momepy/functional/tests/test_diversity.py
@@ -489,6 +489,18 @@ def test_describe_agg(self):
         assert_result(df["sum"], expected_fl_area_sum, self.df_streets)
         assert_result(df["mean"], expected_fl_area_mean, self.df_streets)
 
+    @pytest.mark.skipif(
+        not PD_210, reason="aggregation is different in previous pandas versions"
+    )
+    def test_describe_cols(self):
+        df = mm.describe_agg(
+            self.df_buildings["area"],
+            self.df_buildings["nID"],
+            self.df_streets.index,
+            statistics=["min", "max"],
+        )
+        assert list(df.columns) == ["min", "max"]
+
     @pytest.mark.skipif(
         not PD_210, reason="aggregation is different in previous pandas versions"
     )
@@ -539,6 +551,19 @@ def test_describe_reached_input_equality(self):
             island_result_df.values, island_result_ndarray.values, equal_nan=True
         )
 
+    @pytest.mark.skipif(
+        not PD_210, reason="aggregation is different in previous pandas versions"
+    )
+    def test_describe_reached_cols(self):
+        df = mm.describe_reached_agg(
+            self.df_buildings["fl_area"],
+            self.df_buildings["nID"],
+            graph=self.graph_sw,
+            q=(10, 90),
+            statistics=["min", "max"],
+        )
+        assert list(df.columns) == ["min", "max"]
+
     @pytest.mark.skipif(
         not PD_210, reason="aggregation is different in previous pandas versions"
     )