1# -*- coding: utf-8 -*-
2# Copyright (c) 2017-present, Facebook, Inc.
3# All rights reserved.
4#
5# This source code is licensed under the BSD-style license found in the
6# LICENSE file in the root directory of this source tree. An additional grant
7# of patent rights can be found in the PATENTS file in the same directory.
8
9from __future__ import absolute_import, division, print_function
10
11import warnings
12
13import numpy as np
14import pandas as pd
15
16import fbprophet.hdays as hdays_part2
17import holidays as hdays_part1
18
19
20def get_holiday_names(country):
21    """Return all possible holiday names of given country
22
23    Parameters
24    ----------
25    country: country name
26
27    Returns
28    -------
29    A set of all possible holiday names of given country
30    """
31    years = np.arange(1995, 2045)
32    try:
33        with warnings.catch_warnings():
34            warnings.simplefilter("ignore")
35            holiday_names = getattr(hdays_part2, country)(years=years).values()
36    except AttributeError:
37        try:
38            holiday_names = getattr(hdays_part1, country)(years=years).values()
39        except AttributeError:
40            raise AttributeError(
41                "Holidays in {} are not currently supported!".format(country))
42    return set(holiday_names)
43
44
45def make_holidays_df(year_list, country):
46    """Make dataframe of holidays for given years and countries
47
48    Parameters
49    ----------
50    year_list: a list of years
51    country: country name
52
53    Returns
54    -------
55    Dataframe with 'ds' and 'holiday', which can directly feed
56    to 'holidays' params in Prophet
57    """
58    try:
59        holidays = getattr(hdays_part2, country)(years=year_list)
60    except AttributeError:
61        try:
62            holidays = getattr(hdays_part1, country)(years=year_list)
63        except AttributeError:
64            raise AttributeError(
65                "Holidays in {} are not currently supported!".format(country))
66    holidays_df = pd.DataFrame(list(holidays.items()), columns=['ds', 'holiday'])
67    holidays_df.reset_index(inplace=True, drop=True)
68    holidays_df['ds'] = pd.to_datetime(holidays_df['ds'])
69    return (holidays_df)
70