Make cpi a little bit more robust

This commit is contained in:
jeroen
2023-02-28 11:16:29 +01:00
parent f8fcd174a9
commit d2a388baf4
3 changed files with 25 additions and 12 deletions

View File

@@ -1096,7 +1096,7 @@ class EconomyController(BaseController):
parser.add_argument(
"--no-smart-select",
action="store_true",
action="store_false",
dest="smart_select",
help="Whether to assist with selection",
default=True,

View File

@@ -28,7 +28,6 @@ CPI_COUNTRIES = [
"austria",
"belgium",
"brazil",
"bulgaria",
"canada",
"chile",
"china",
@@ -316,6 +315,7 @@ def get_cpi(
frequency: str = "",
harmonized: bool = False,
smart_select: bool = True,
options: bool = False,
start_date: Optional[str] = None,
end_date: Optional[str] = None,
) -> Tuple[pd.DataFrame, dict]:
@@ -333,14 +333,18 @@ def get_cpi(
Whether you wish to obtain harmonized data.
smart_select: bool
Whether to assist with the selection.
options: bool
Whether to return the options.
start_date: Optional[str]
Start date, formatted YYYY-MM-DD
end_date: Optional[str]
End date, formatted YYYY-MM-DD
"""
series = pd.read_csv(harmonized_cpi_path) if harmonized else pd.read_csv(cpi_path)
if options:
return series.drop(["series_id"], axis=1)
step_1 = series[series["country"].str.contains("|".join(countries))]
step_2 = step_1[step_1["units"] == units]
step_3 = step_2[step_2["frequency"] == frequency]
@@ -363,7 +367,11 @@ def get_cpi(
series = step_3
if series.empty:
console.print("The combination of parameters does not result in any data.")
console.print(
"The combination of parameters does not result in any data. Please consider "
"using the `options` parameter to see the available options. Note that there "
"are two options list, one with `harmonized` and one without."
)
return pd.DataFrame()
series_dictionary = {}

View File

@@ -222,19 +222,20 @@ def plot_cpi(
else pd.read_csv(fred_model.cpi_path)
)
if options:
return print_rich_table(series.drop(["series_id"], axis=1))
df = fred_model.get_cpi(
countries=countries,
units=units,
frequency=frequency,
harmonized=harmonized,
smart_select=smart_select,
options=options,
start_date=start_date,
end_date=end_date,
)
if options:
return print_rich_table(series.drop(["series_id"], axis=1))
# This plot has 1 axis
if not external_axes:
_, ax = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI)
@@ -261,11 +262,15 @@ def plot_cpi(
ax.set_title(title, fontsize=15)
ax.legend(prop={"size": 8})
else:
country = countries[0].replace("_", " ").title()
title = (
f"{'Harmonized ' if harmonized else ''}Consumer Price Index for {country}"
)
ax.set_title(title)
country = [
country.replace("_", " ").title()
for country in countries
if df.columns.str.contains(country)
]
if country:
title = f"{'Harmonized ' if harmonized else ''}Consumer Price Index for {country[0]}"
ax.set_title(title)
if units == "growth_same":
ax.set_ylabel("Growth Same Period (%)")