fix(response): assert expected columns are present even on empty output (#165)

Author: SlowMo24

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Fixed
+
+- Assert that the expected columns are also present if the result is empty
+
 ## 0.3.2
 
 ### Fixed

ohsome/response.py

@@ -4,7 +4,7 @@
 """Class for ohsome API response"""
 
 import json
-from typing import Optional
+from typing import Optional, Union
 
 import geopandas as gpd
 import pandas as pd
@@ -25,7 +25,7 @@ def __init__(self, response=None, url=None, params=None):
 
     def as_dataframe(
         self, multi_index: Optional[bool] = True, explode_tags: Optional[tuple] = ()
-    ):
+    ) -> Union[pd.DataFrame, gpd.GeoDataFrame]:
         """
         Converts the ohsome response to a pandas.DataFrame or a geopandas.GeoDataFrame if the
         response contains geometries
@@ -40,13 +40,7 @@ def as_dataframe(
         else:
             return self._as_geodataframe(multi_index, explode_tags)
 
-    def _as_dataframe(self, multi_index=True):
-        """
-        Converts the ohsome response to a pandas.DataFrame
-        :param multi_index: If true returns the dataframe with a multi index
-        :return: pandas.DataFrame
-        """
-
+    def _as_dataframe(self, multi_index=True) -> pd.DataFrame:
         groupby_names = []
         if "result" in self.data.keys():
             result_df = pd.DataFrame().from_records(self.data["result"])
@@ -79,15 +73,17 @@ def _as_dataframe(self, multi_index=True):
 
     def _as_geodataframe(
         self, multi_index: Optional[bool] = True, explode_tags: Optional[tuple] = ()
-    ):
-        """
-        Converts the ohsome response to a geopandas.GeoDataFrame
-        :param multi_index: If true returns the dataframe with a multi index
-        :return: geopandas.GeoDataFrame
-        """
-
+    ) -> gpd.GeoDataFrame:
         if len(self.data["features"]) == 0:
-            return gpd.GeoDataFrame(crs="epsg:4326", columns=["@osmId", "geometry"])
+            return gpd.GeoDataFrame(
+                crs="epsg:4326",
+                columns=["@osmId", "geometry"]
+                + (
+                    list(explode_tags) + ["@other_tags"]
+                    if explode_tags is not None
+                    else []
+                ),
+            )
 
         try:
             if explode_tags is not None:
@@ -128,7 +124,7 @@ def _as_geodataframe(
 
         return features.sort_index()
 
-    def to_json(self, outfile):
+    def to_json(self, outfile) -> None:
         """
         Write response to json file
         :return:
@@ -137,7 +133,7 @@ def to_json(self, outfile):
         with open(outfile, "w", encoding="utf-8") as dst:
             json.dump(self.data, dst, indent=2, ensure_ascii=False)
 
-    def _set_index(self, result_df, groupby_names):
+    def _set_index(self, result_df, groupby_names) -> None:
         """
         Set multi-index based on groupby names and time
         :param result_df:

ohsome/test/test_response.py

@@ -8,8 +8,11 @@
 import pandas as pd
 import pytest
 from geopandas.testing import assert_geodataframe_equal
+from requests import Response
 from shapely import Point
 
+from ohsome import OhsomeResponse
+
 
 @pytest.mark.vcr
 def test_elements_count(base_client):
@@ -406,6 +409,7 @@ def test_empty_geodataframe(base_client):
 
     assert isinstance(result, gpd.GeoDataFrame)
     assert result.empty
+    assert result.columns.to_list() == ["@osmId", "geometry", "@other_tags"]
 
 
 @pytest.mark.vcr
@@ -543,3 +547,22 @@ def test_explode_tags_missing_in_response(dummy_ohsome_response):
     )
 
     assert_geodataframe_equal(computed_df, expected_df, check_like=True)
+
+
+def test_explode_tags_present_on_empty_result():
+    """Test if exploded tags are present in an empty results."""
+    expected_df = gpd.GeoDataFrame(
+        columns=["@osmId", "some_key", "some_other_key", "@other_tags", "geometry"],
+        crs="EPSG:4326",
+    )
+
+    response = Response()
+    response._content = (
+        '{"attribution":{"url":"https://ohsome.org/copyrights","text":"© OpenStreetMap contributors"},'
+        '"apiVersion":"1.10.1","type":"FeatureCollection","features":[]}'
+    ).encode()
+    computed_df = OhsomeResponse(response=response).as_dataframe(
+        explode_tags=("some_key", "some_other_key")
+    )
+
+    assert_geodataframe_equal(computed_df, expected_df, check_like=True)