Skip to content

Commit

Permalink
add more countries
Browse files Browse the repository at this point in the history
  • Loading branch information
deeleeramone committed May 15, 2024
1 parent ef42545 commit cd54a51
Show file tree
Hide file tree
Showing 3 changed files with 96 additions and 13 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -1334,7 +1334,7 @@ def fixedincome_government_yield_curve( # noqa: PLR0912
colors = LARGE_CYCLER
color_count = 0

figure = OpenBBFigure().create_subplots(shared_xaxes=False)
figure = OpenBBFigure().create_subplots(shared_xaxes=True)
figure.update_layout(ChartStyle().plotly_template.get("layout", {}))

def create_fig(figure, df, dates, color_count, country: Optional[str] = None):
Expand Down Expand Up @@ -1417,6 +1417,18 @@ def create_fig(figure, df, dates, color_count, country: Optional[str] = None):
title="Maturity",
ticklen=0,
showgrid=False,
type="category",
categoryorder="array",
categoryarray=(
[
(
d.split("_")[1] + " " + d.split("_")[0].title()
if d != "long_term"
else "Long Term"
)
for d in maturities
]
),
),
yaxis=dict(
title="Yield (%)",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from openbb_core.provider.utils.descriptions import QUERY_DESCRIPTIONS
from openbb_core.provider.utils.errors import EmptyDataError
from openbb_econdb.utils import helpers
from openbb_econdb.utils.yield_curves import COUNTRIES, DAILY, DAILY_COUNTRIES
from openbb_econdb.utils.yield_curves import COUNTRIES, COUNTRIES_DICT, COUNTRIES_LIST
from pandas import Categorical, DataFrame, DatetimeIndex
from pydantic import Field, field_validator

Expand All @@ -27,8 +27,10 @@ class EconDbYieldCurveQueryParams(YieldCurveQueryParams):
}
country: COUNTRIES = Field(
default="united_states",
description=QUERY_DESCRIPTIONS.get("country", ""),
json_schema_extra={"choices": DAILY_COUNTRIES},
description=QUERY_DESCRIPTIONS.get("country", "")
+ " New Zealand, Mexico, Singapore, and Thailand have only monthly data."
+ " The nearest date to the requested one will be used.",
json_schema_extra={"choices": COUNTRIES_LIST},
)
use_cache: bool = Field(
default=True,
Expand All @@ -41,8 +43,8 @@ def validate_country(cls, v) -> str:
"""Validate the country."""
if v is None:
return "united_states"
if v not in DAILY_COUNTRIES:
raise ValueError(f"Country must be one of {DAILY_COUNTRIES}")
if v not in COUNTRIES_DICT:
raise ValueError(f"Country must be one of {COUNTRIES_DICT}")
return v


Expand Down Expand Up @@ -77,7 +79,7 @@ async def aextract_data( # pylint: disable=R0914.R0912,R0915
token = await helpers.create_token(use_cache=query.use_cache)
credentials.update({"econdb_api_key": token}) # type: ignore
base_url = "https://www.econdb.com/api/series/?ticker="
symbols = list(DAILY[query.country].keys())
symbols = list(COUNTRIES_DICT[query.country].keys())
url = (
base_url
+ f"%5B{','.join(symbols)}%5D&page_size=50&format=json&token={token}"
Expand Down Expand Up @@ -117,15 +119,15 @@ def transform_data(
**kwargs: Any,
) -> AnnotatedResult[List[EconDbYieldCurveData]]:
"""Transform the data."""
maturity_order = list(DAILY[query.country].values())
maturity_order = list(COUNTRIES_DICT[query.country].values())
dates = query.date.split(",") # type: ignore
dates_list = DatetimeIndex(dates)
new_data: Dict = {}
metadata: Dict = {}
# Unpack the data for each maturity.
for item in data:
ticker = item.get("ticker")
maturity = DAILY[query.country].get(ticker)
maturity = COUNTRIES_DICT[query.country].get(ticker)
dataset = item.get("dataset")
additional_metadata = item.get("additional_metadata")
units = "percent"
Expand Down Expand Up @@ -171,8 +173,13 @@ def transform_data(
by=["date", "maturity"]
).reset_index(drop=True)
flattened_data.loc[:, "date"] = flattened_data["date"].dt.strftime("%Y-%m-%d")
records = flattened_data.to_dict(orient="records")
records = flattened_data.copy()
records["rate"] = records["rate"].fillna("N/A").replace("N/A", None)
return AnnotatedResult(
result=[EconDbYieldCurveData.model_validate(r) for r in records],
result=[
EconDbYieldCurveData.model_validate(r)
for r in records.to_dict("records")
if r.get("rate")
],
metadata=metadata,
)
Original file line number Diff line number Diff line change
Expand Up @@ -135,21 +135,85 @@
"FRB_H15.F5280.D.US": "year_30",
},
}

MONTHLY = {
"thailand": {
"BOTFM_RT_001_S2.36.M.TH": "month_1",
"BOTFM_RT_001_S2.37.M.TH": "month_3",
"BOTFM_RT_001_S2.38.M.TH": "month_6",
"BOTFM_RT_001_S2.39.M.TH": "year_1",
"BOTFM_RT_001_S2.40.M.TH": "year_2",
"BOTFM_RT_001_S2.41.M.TH": "year_3",
"BOTFM_RT_001_S2.42.M.TH": "year_4",
"BOTFM_RT_001_S2.43.M.TH": "year_5",
"BOTFM_RT_001_S2.44.M.TH": "year_6",
"BOTFM_RT_001_S2.45.M.TH": "year_7",
"BOTFM_RT_001_S2.46.M.TH": "year_8",
"BOTFM_RT_001_S2.47.M.TH": "year_9",
"BOTFM_RT_001_S2.48.M.TH": "year_10",
"BOTFM_RT_001_S2.49.M.TH": "year_11",
"BOTFM_RT_001_S2.50.M.TH": "year_12",
"BOTFM_RT_001_S2.51.M.TH": "year_13",
"BOTFM_RT_001_S2.52.M.TH": "year_14",
"BOTFM_RT_001_S2.53.M.TH": "year_15",
"BOTFM_RT_001_S2.54.M.TH": "year_16",
"BOTFM_RT_001_S2.55.M.TH": "year_17",
"BOTFM_RT_001_S2.56.M.TH": "year_18",
"BOTFM_RT_001_S2.57.M.TH": "year_19",
"BOTFM_RT_001_S2.58.M.TH": "year_20",
},
"new_zealand": {
"RBNZ_IR.MB01N.M.NZ": "month_1",
"RBNZ_IR.MB02N.M.NZ": "month_2",
"RBNZ_IR.MB03N.M.NZ": "month_3",
"RBNZ_IR.MG101N.M.NZ": "year_1",
"RBNZ_IR.MG102N.M.NZ": "year_2",
"RBNZ_IR.MG105N.M.NZ": "year_5",
"RBNZ_IR.MG110N.M.NZ": "year_10",
},
"mexico": {
"BXC_PS.SF4830.M.MX": "month_1",
"BXC_PS.SF4831.M.MX": "month_3",
"BXC_PS.SF4832.M.MX": "month_6",
"BXC_PS.SF4833.M.MX": "year_1",
"BXC_PS.SF17990.M.MX": "year_3",
"BXC_PS.SF98793.M.MX": "year_5",
"BXC_PS.SF33229.M.MX": "year_7",
"BXC_PS.SF30057.M.MX": "year_10",
"BXC_PS.SF41841.M.MX": "year_20",
"BXC_PS.SF60719.M.MX": "year_30",
},
"singapore": {
"SSG_IR.M70007118.M.SG": "year_1",
"SSG_IR.M70007113.M.SG": "year_2",
"SSG_IR.M70007110.M.SG": "year_5",
"SSG_IR.M70007115.M.SG": "year_10",
"SSG_IR.M70007119.M.SG": "year_15",
"SSG_IR.M70007120.M.SG": "year_20",
},
}
DAILY_COUNTRIES = list(DAILY.keys())

MONTHLY_COUNTRIES = list(MONTHLY.keys())
COUNTRIES_LIST = sorted(DAILY_COUNTRIES + MONTHLY_COUNTRIES)
COUNTRIES_DICT = {
**DAILY,
**MONTHLY,
}
COUNTRIES = Literal[
"australia",
"canada",
"china",
"hong_kong",
"india",
"japan",
"mexico",
"new_zealand",
"russia",
"saudi_arabia",
"singapore",
"south_africa",
"south_korea",
"taiwan",
"thailand",
"united_kingdom",
"united_states",
]

0 comments on commit cd54a51

Please sign in to comment.