From a47d59a75d299a5304fe989363da6bf1521075a5 Mon Sep 17 00:00:00 2001 From: Moritz Schott Date: Wed, 7 Aug 2024 18:31:07 +0200 Subject: [PATCH] complete typing --- ohsome/response.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/ohsome/response.py b/ohsome/response.py index 5ed5f85..b7c76e1 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,7 +40,7 @@ def as_dataframe( else: return self._as_geodataframe(multi_index, explode_tags) - def _as_dataframe(self, multi_index=True): + 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"]) @@ -73,7 +73,7 @@ def _as_dataframe(self, multi_index=True): def _as_geodataframe( self, multi_index: Optional[bool] = True, explode_tags: Optional[tuple] = () - ): + ) -> gpd.GeoDataFrame: if len(self.data["features"]) == 0: return gpd.GeoDataFrame( crs="epsg:4326", @@ -124,7 +124,7 @@ def _as_geodataframe( return features.sort_index() - def to_json(self, outfile): + def to_json(self, outfile) -> dict: """ Write response to json file :return: @@ -133,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: