Merge pull request #132 from lahdjirayhan/improve-get-forecast

fix: Improve `get_city_forecast` to pass tests

src/ozone/ozone.py

@@ -249,7 +249,12 @@ def _extract_forecast_data(self, data_obj: Any) -> pandas.DataFrame:
  dict_of_frames[pol] = pandas.DataFrame(lst).set_index("day")
 
  df = pandas.concat(dict_of_frames, axis=1)
+
+ # Convert to numeric while making non-numerics nan,
+ # and then convert to float, just in case there's int
+ df = df.apply(lambda x: pandas.to_numeric(x, errors="coerce")).astype(float)
  df.index = pandas.to_datetime(df.index)
+ df = df.reset_index().rename(columns={"day": "date"})
  return df
 
  def _check_and_get_data_obj(

tests/test_get_city_forecast.py

@@ -20,20 +20,17 @@ def test_return_value_and_format():
 
 
 @pytest.mark.vcr
-def test_index_and_column_types():
+def test_column_types():
  result = api.get_city_forecast("london")
 
- # The index contains date information
- assert pd_types.is_datetime64_any_dtype(result.index)
+ # The date column contains date information
+ assert pd_types.is_datetime64_any_dtype(result["date"])
 
  # Check that all pollutants and statistics are in column
+ # and are of float type
  for param in ["pm25", "pm10", "o3", "uvi"]:
  for stat in ["min", "max", "avg"]:
- assert (param, stat) in result
-
- # Check that all columns are of float type
- for col in result:
- assert pd_types.is_float_dtype(result[col])
+ assert pd_types.is_float_dtype(result[(param, stat)])
 
 
 @pytest.mark.vcr