diff --git a/CHANGELOG.md b/CHANGELOG.md index 11c697c..15ff609 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/ohsome/response.py b/ohsome/response.py index c6880b1..4301475 100644 --- a/ohsome/response.py +++ b/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: diff --git a/ohsome/test/test_response.py b/ohsome/test/test_response.py index e1514d4..fa3c3c5 100644 --- a/ohsome/test/test_response.py +++ b/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)