prophet/python/fbprophet/make_holidays.py
Greg Rafferty 3d0bb6e081
fix for comma-separated holidays (#1638)
* use holidays.get_list() to return list of holidays

* make_holidays_df now accepts states
2020-08-27 17:49:19 -07:00

68 lines
2.1 KiB
Python

# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import absolute_import, division, print_function
import warnings
import numpy as np
import pandas as pd
import fbprophet.hdays as hdays_part2
import holidays as hdays_part1
def get_holiday_names(country):
"""Return all possible holiday names of given country
Parameters
----------
country: country name
Returns
-------
A set of all possible holiday names of given country
"""
years = np.arange(1995, 2045)
try:
with warnings.catch_warnings():
warnings.simplefilter("ignore")
holiday_names = getattr(hdays_part2, country)(years=years).values()
except AttributeError:
try:
holiday_names = getattr(hdays_part1, country)(years=years).values()
except AttributeError as e:
raise AttributeError(
"Holidays in {} are not currently supported!".format(country)) from e
return set(holiday_names)
def make_holidays_df(year_list, country, province=None, state=None):
"""Make dataframe of holidays for given years and countries
Parameters
----------
year_list: a list of years
country: country name
Returns
-------
Dataframe with 'ds' and 'holiday', which can directly feed
to 'holidays' params in Prophet
"""
try:
holidays = getattr(hdays_part2, country)(years=year_list)
except AttributeError:
try:
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([(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)