diff --git a/python/fbprophet/make_holidays.py b/python/fbprophet/make_holidays.py index 27b3c49..fc52dc0 100644 --- a/python/fbprophet/make_holidays.py +++ b/python/fbprophet/make_holidays.py @@ -40,7 +40,7 @@ def get_holiday_names(country): return set(holiday_names) -def make_holidays_df(year_list, country, province=None): +def make_holidays_df(year_list, country, province=None, state=None): """Make dataframe of holidays for given years and countries Parameters @@ -57,11 +57,12 @@ def make_holidays_df(year_list, country, province=None): holidays = getattr(hdays_part2, country)(years=year_list) except AttributeError: try: - holidays = getattr(hdays_part1, country)(prov=province,years=year_list) + holidays = getattr(hdays_part1, country)(prov=province, state=state, years=year_list) except AttributeError as e: raise AttributeError( "Holidays in {} are not currently supported!".format(country)) from e - holidays_df = pd.DataFrame(list(holidays.items()), columns=['ds', 'holiday']) + holidays_df = pd.DataFrame([(date, holidays.get_list(date)) for date in holidays], columns=['ds', 'holiday']) + holidays_df = holidays_df.explode('holiday') holidays_df.reset_index(inplace=True, drop=True) holidays_df['ds'] = pd.to_datetime(holidays_df['ds']) return (holidays_df)