{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-20T02:46:08.143255Z",
     "start_time": "2019-02-20T02:46:08.086340Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import collections\n",
    "import datetime\n",
    "import os\n",
    "import inspect\n",
    "\n",
    "def _as_np_date(val):\n",
    "    '''\n",
    "    Convert a pandas timestamp, string, np.datetime64('M8[ns]'), datetime, date to a numpy datetime64 [D] and remove time info.\n",
    "    Returns None if the value cannot be converted\n",
    "    >>> _as_np_date(pd.Timestamp('2016-05-01 3:55:00'))\n",
    "    numpy.datetime64('2016-05-01')\n",
    "    >>> _as_np_date('2016-05-01')\n",
    "    numpy.datetime64('2016-05-01')\n",
    "    >>> x = pd.DataFrame({'x' : [np.datetime64('2015-01-01 05:00:00'), np.datetime64('2015-02-01 06:00:00')]})\n",
    "    >>> _as_np_date(x.x)\n",
    "    array(['2015-01-01', '2015-02-01'], dtype='datetime64[D]')\n",
    "    >>> _as_np_date(pd.Series([np.datetime64('2015-01-01 05:00:00'), np.datetime64('2015-02-01 06:00:00')]))\n",
    "    array(['2015-01-01', '2015-02-01'], dtype='datetime64[D]')\n",
    "    >>> x = pd.DataFrame({'x' : [1, 2]}, index = [np.datetime64('2015-01-01 05:00:00'), np.datetime64('2015-02-01 06:00:00')])\n",
    "    >>> _as_np_date(x.index)\n",
    "    array(['2015-01-01', '2015-02-01'], dtype='datetime64[D]')\n",
    "    '''\n",
    "    if isinstance(val, np.datetime64): \n",
    "        return val.astype('M8[D]')\n",
    "    if isinstance(val, str) or isinstance(val, datetime.date) or isinstance(val, datetime.datetime):\n",
    "        np_date = np.datetime64(val).astype('M8[D]')\n",
    "        if isinstance(np_date.astype(datetime.datetime), int): # User can pass in a string like 20180101 which gets parsed as a year\n",
    "            raise Exception(f'invalid date: {val}')\n",
    "        return np_date\n",
    "    if isinstance(val, pd.Timestamp): \n",
    "        return timestamp.to_datetime64().astype('M8[D]')\n",
    "    if isinstance(val, pd.Series) or isinstance(val, pd.DatetimeIndex):\n",
    "        return val.values.astype('M8[D]')\n",
    "    if isinstance(val, np.ndarray) and val.dtype == np.dtype('M8[ns]'): \n",
    "        return val.astype('M8[D]')\n",
    "    return None\n",
    "\n",
    "\n",
    "def _normalize(start, end, include_first, include_last):\n",
    "    '''\n",
    "    Given a start and end date, return a new start and end date, taking into account include_first and include_last flags\n",
    "    \n",
    "    Args:\n",
    "        start: start date, can be string or datetime or np.datetime64\n",
    "        end: end_date, can be string or datetime or np.datetime64\n",
    "        include_first (bool): whether to increment start date by 1\n",
    "        include_last (bool): whether to increment end date by 1\n",
    "\n",
    "    Return:\n",
    "        np.datetime64, np.datetime64 : new start and end dates\n",
    "        \n",
    "    >>> x = np.arange(np.datetime64('2017-01-01'), np.datetime64('2017-01-10'))\n",
    "    >>> y = x.astype('M8[ns]')\n",
    "    >>> xcp = x.copy()\n",
    "    >>> ycp = y.copy()\n",
    "    >>> _normalize(x, y, False, False) # doctest:+ELLIPSIS +NORMALIZE_WHITESPACE\n",
    "    (array(...\n",
    "    >>> (x == xcp).all()\n",
    "    True\n",
    "    >>> (y == ycp).all()\n",
    "    True\n",
    "    '''\n",
    "    s = _as_np_date(start)\n",
    "    e = _as_np_date(end)\n",
    "    \n",
    "    if s is None: s = start.copy()\n",
    "    if e is None: e = end.copy()\n",
    "\n",
    "    if not include_first:\n",
    "        s += np.timedelta64(1, 'D')\n",
    "\n",
    "    if include_last:\n",
    "        e += np.timedelta64(1, 'D')\n",
    "\n",
    "    return s,e\n",
    "\n",
    "def read_holidays(calendar_name, dirname = None):\n",
    "    '''\n",
    "    Reads a csv with a holidays column containing holidays (not including weekends)\n",
    "    '''\n",
    "    if dirname is None: dirname = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))\n",
    "    \n",
    "    if not os.path.isdir(dirname + '/refdata'):\n",
    "        if os.path.isdir(dirname + '/../refdata'):\n",
    "            dirname = dirname + '/../'\n",
    "        else:\n",
    "            raise Exception(f'path {dirname}/refdata and {dirname}/../refdata do not exist')\n",
    "        #raise Exception(f'path {dirname}/refdata does not exist')\n",
    "    df = pd.read_csv(f'{dirname}/refdata/holiday_calendars/{calendar_name}.csv')\n",
    "    holidays = pd.to_datetime(df.holidays, format='%Y-%m-%d').values.astype('M8[D]')\n",
    "    return holidays\n",
    "\n",
    "class Calendar(object):\n",
    "    \n",
    "    NYSE = \"nyse\"\n",
    "    EUREX = \"eurex\"\n",
    "    _calendars = {}\n",
    "    \n",
    "    def __init__(self, holidays):\n",
    "        '''\n",
    "        Do not use this function directly.  Use Calendar.get_calendar instead\n",
    "        Args:\n",
    "            holidays (np.array of datetime64[D]): holidays for this calendar, excluding weekends\n",
    "        '''\n",
    "        self.bus_day_cal = np.busdaycalendar(holidays = holidays)\n",
    "        \n",
    "    def is_trading_day(self, dates):\n",
    "        '''\n",
    "        Returns whether the date is not a holiday or a weekend\n",
    "        \n",
    "        Args:\n",
    "            dates: str or datetime.datetime or np.datetime64[D] or numpy array of np.datetime64[D]\n",
    "        Return:\n",
    "            bool: Whether this date is a trading day\n",
    "        \n",
    "        >>> import datetime\n",
    "        >>> eurex = Calendar.get_calendar(Calendar.EUREX)\n",
    "        >>> eurex.is_trading_day('2016-12-25')\n",
    "        False\n",
    "        >>> eurex.is_trading_day(datetime.date(2016, 12, 22))\n",
    "        True\n",
    "        >>> nyse = Calendar.get_calendar(Calendar.NYSE)\n",
    "        >>> nyse.is_trading_day('2017-04-01') # Weekend\n",
    "        False\n",
    "        >>> nyse.is_trading_day(np.arange('2017-04-01', '2017-04-09', dtype = np.datetime64)) # doctest:+ELLIPSIS\n",
    "        array([False, False,  True,  True,  True,  True,  True, False]...)\n",
    "        '''\n",
    "        if isinstance(dates, str) or isinstance(dates, datetime.date): \n",
    "            dates = np.datetime64(dates, 'D')\n",
    "            if isinstance(dates.astype(datetime.datetime), int): # User can pass in a string like 20180101 which gets parsed as a year\n",
    "                raise Exception(f'invalid date: {dates}')\n",
    "        if isinstance(dates, pd.Series): dates = dates.values\n",
    "        return np.is_busday(dates.astype('M8[D]'), busdaycal = self.bus_day_cal)\n",
    "    \n",
    "    def num_trading_days(self, start, end, include_first = False, include_last = True):\n",
    "        '''\n",
    "        Count the number of trading days between two date series including those two dates\n",
    "        You can pass in a string like '2009-01-01' or a python date or a pandas series for \n",
    "        start and end\n",
    "        \n",
    "        >>> eurex = Calendar.get_calendar(Calendar.EUREX)\n",
    "        >>> eurex.num_trading_days('2009-01-01', '2011-12-31')\n",
    "        772.0\n",
    "        >>> dates = pd.date_range('20130101',periods=8)\n",
    "        >>> increments = np.array([5, 0, 3, 9, 4, 10, 15, 29])\n",
    "        >>> import warnings\n",
    "        >>> import pandas as pd\n",
    "        >>> warnings.filterwarnings(action = 'ignore', category = pd.errors.PerformanceWarning)\n",
    "        >>> dates2 = dates + increments * dates.freq\n",
    "        >>> df = pd.DataFrame({'x': dates, 'y' : dates2})\n",
    "        >>> df.iloc[4]['x'] = np.nan\n",
    "        >>> df.iloc[6]['y'] = np.nan\n",
    "        >>> nyse = Calendar.get_calendar(Calendar.NYSE)\n",
    "        >>> np.set_printoptions(formatter = {'float' : lambda x : f'{x:.1f}'})  # After numpy 1.13 positive floats don't have a leading space for sign\n",
    "        >>> print(nyse.num_trading_days(df.x, df.y))\n",
    "        [3.0 0.0 1.0 5.0 nan 8.0 nan 20.0]\n",
    "        '''\n",
    "        s_tmp, e_tmp = _normalize(start, end, include_first, include_last)\n",
    "        # np.busday_count does not like nat dates so convert them to known numbers and then check for them.\n",
    "        DUMMY_BEGIN_DATE = np.datetime64('1700-01-01')\n",
    "        DUMMY_END_DATE = np.datetime64('1829-03-08')\n",
    "        DUMMY_DIFF = 33701.\n",
    "        s = np.where((np.isnat(s_tmp) | np.isnat(e_tmp)), DUMMY_BEGIN_DATE, s_tmp)\n",
    "        e = np.where((np.isnat(s_tmp) | np.isnat(e_tmp)), DUMMY_END_DATE, e_tmp)\n",
    "        count = np.busday_count(s, e, busdaycal = self.bus_day_cal).astype(np.float)\n",
    "        count = np.where(count == DUMMY_DIFF, np.nan, count)\n",
    "        if isinstance(start, str) and isinstance(end, str):\n",
    "            return np.ndarray.item(count)\n",
    "        if not isinstance(start, collections.abc.Iterable) and not isinstance(end, collections.abc.Iterable):\n",
    "            return np.ndarray.item(count)\n",
    "        return count\n",
    "        \n",
    "    def get_trading_days(self, start, end, include_first = False, include_last = True):\n",
    "        '''\n",
    "        Get back a list of numpy dates that are trading days between the start and end\n",
    "        \n",
    "        >>> nyse = Calendar.get_calendar(Calendar.NYSE)\n",
    "        >>> nyse.get_trading_days('2005-01-01', '2005-01-08')\n",
    "        array(['2005-01-03', '2005-01-04', '2005-01-05', '2005-01-06', '2005-01-07'], dtype='datetime64[D]')\n",
    "        >>> nyse.get_trading_days(datetime.date(2005, 1, 1), datetime.date(2005, 2, 1))\n",
    "        array(['2005-01-03', '2005-01-04', '2005-01-05', '2005-01-06',\n",
    "               '2005-01-07', '2005-01-10', '2005-01-11', '2005-01-12',\n",
    "               '2005-01-13', '2005-01-14', '2005-01-18', '2005-01-19',\n",
    "               '2005-01-20', '2005-01-21', '2005-01-24', '2005-01-25',\n",
    "               '2005-01-26', '2005-01-27', '2005-01-28', '2005-01-31', '2005-02-01'], dtype='datetime64[D]')\n",
    "        >>> nyse.get_trading_days(datetime.date(2016, 1, 5), datetime.date(2016, 1, 29), include_last = False)\n",
    "        array(['2016-01-06', '2016-01-07', '2016-01-08', '2016-01-11',\n",
    "               '2016-01-12', '2016-01-13', '2016-01-14', '2016-01-15',\n",
    "               '2016-01-19', '2016-01-20', '2016-01-21', '2016-01-22',\n",
    "               '2016-01-25', '2016-01-26', '2016-01-27', '2016-01-28'], dtype='datetime64[D]')\n",
    "        >>> nyse.get_trading_days('2017-07-04', '2017-07-08', include_first = False)\n",
    "        array(['2017-07-05', '2017-07-06', '2017-07-07'], dtype='datetime64[D]')\n",
    "        >>> nyse.get_trading_days(np.datetime64('2017-07-04'), np.datetime64('2017-07-08'), include_first = False)\n",
    "        array(['2017-07-05', '2017-07-06', '2017-07-07'], dtype='datetime64[D]')\n",
    "        ''' \n",
    "        s, e = _normalize(start, end, include_first, include_last)\n",
    "        dates = np.arange(s, e, dtype='datetime64[D]')\n",
    "        dates = dates[np.is_busday(dates, busdaycal=self.bus_day_cal)]\n",
    "        return dates\n",
    "    \n",
    "    def add_trading_days(self, start, num_days, roll = 'raise'):\n",
    "        '''\n",
    "        Adds trading days to a start date\n",
    "        \n",
    "        Args:\n",
    "            start: np.datetime64 or str or datetime\n",
    "            num_days (int): number of trading days to add\n",
    "            roll (str, optional): one of 'raise', 'nat', 'forward', 'following', 'backward', 'preceding', 'modifiedfollowing', 'modifiedpreceding'}\n",
    "                From numpy documentation: \n",
    "                How to treat dates that do not fall on a valid day. The default is ‘raise’.\n",
    "                'raise' means to raise an exception for an invalid day.\n",
    "                'nat' means to return a NaT (not-a-time) for an invalid day.\n",
    "                'forward' and 'following’ mean to take the first valid day later in time.\n",
    "                'backward' and 'preceding' mean to take the first valid day earlier in time.\n",
    "                'modifiedfollowing' means to take the first valid day later in time unless it is across a Month boundary, \n",
    "                in which case to take the first valid day earlier in time.\n",
    "                'modifiedpreceding' means to take the first valid day earlier in time unless it is across a Month boundary, \n",
    "                in which case to take the first valid day later in time.\n",
    "        \n",
    "        Return:\n",
    "            np.datetime64[D]: The date num_days trading days after start\n",
    "            \n",
    "        >>> calendar = Calendar.get_calendar(Calendar.NYSE)\n",
    "        >>> calendar.add_trading_days(datetime.date(2015, 12, 24), 1)\n",
    "        numpy.datetime64('2015-12-28')\n",
    "        >>> calendar.add_trading_days(np.datetime64('2017-04-15'), 0, roll = 'preceding') # 4/14/2017 is a Friday and a holiday\n",
    "        numpy.datetime64('2017-04-13')\n",
    "        >>> calendar.add_trading_days(np.datetime64('2017-04-08'), 0, roll = 'preceding') # 4/7/2017 is a Friday and not a holiday\n",
    "        numpy.datetime64('2017-04-07')\n",
    "        '''\n",
    "        start = _as_np_date(start)\n",
    "        out = np.busday_offset(start, num_days, roll = roll, busdaycal = self.bus_day_cal)\n",
    "        return out\n",
    "        \n",
    "    def add_calendar(exchange_name, holidays):\n",
    "        '''\n",
    "        Add a trading calendar to the class level calendars dict\n",
    "        \n",
    "        Args:\n",
    "            exchange_name (str): Name of the exchange.\n",
    "            holidays (np.array of datetime64[D]): holidays for this exchange, excluding weekends\n",
    "        '''\n",
    "        Calendar._calendars[exchange_name] = Calendar(holidays)\n",
    "        \n",
    "    def get_calendar(exchange_name):\n",
    "        '''\n",
    "        Get a calendar object for the given exchange:\n",
    "        \n",
    "        Args:\n",
    "            exchange_name (str): The exchange for which you want a calendar.  Calendar.NYSE, Calendar.EUREX are predefined.\n",
    "            If you want to add a new calendar, use the add_calendar class level function\n",
    "        \n",
    "        Return:\n",
    "            Calendar: The calendar object\n",
    "        '''\n",
    "        \n",
    "        if exchange_name not in Calendar._calendars:\n",
    "            if exchange_name not in [Calendar.NYSE, Calendar.EUREX]:\n",
    "                raise Exception(f'calendar not found: {exchange}')\n",
    "            holidays = read_holidays(exchange_name)\n",
    "            Calendar.add_calendar(exchange_name, holidays)\n",
    "        return Calendar._calendars[exchange_name]\n",
    "    \n",
    "if __name__ == \"__main__\":\n",
    "    import doctest\n",
    "    doctest.testmod(optionflags = doctest.NORMALIZE_WHITESPACE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
