{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-08T03:28:16.330100Z",
     "start_time": "2021-01-08T03:28:15.721016Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running typecheck\n",
      "Success: no issues found in 1 source file\n",
      "\n",
      "running flake8\n",
      "flake8 success\n"
     ]
    }
   ],
   "source": [
    "%%checkall\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from collections.abc import Iterable\n",
    "import datetime\n",
    "import os\n",
    "import inspect\n",
    "import calendar as cal\n",
    "import dateutil.relativedelta as rd\n",
    "\n",
    "from typing import Optional, Tuple, Union, MutableMapping\n",
    "from types import FrameType\n",
    "\n",
    "\n",
    "def _as_np_date(val: Union[pd.Timestamp, str, np.datetime64, datetime.datetime, datetime.date]) -> Optional[np.datetime64]:\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 val.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 np.issubdtype(val.dtype, np.datetime64): \n",
    "        return val.astype('M8[D]')\n",
    "    return None\n",
    "\n",
    "\n",
    "def _normalize_datetime(val: Union[pd.Timestamp, str, np.datetime64, datetime.datetime, datetime.date]) -> Tuple[np.datetime64, np.timedelta64]:\n",
    "    '''\n",
    "    Break up a datetime into numpy date and numpy timedelta.  \n",
    "    \n",
    "    Args:\n",
    "        val: The datetime to normalize.  Can be an array or a single datetime as a string, pandas timestamp, numpy datetime \n",
    "            or python date or datetime\n",
    "    \n",
    "    >>> print(_normalize_datetime(pd.Timestamp('2016-05-01 3:55:00')))\n",
    "    (numpy.datetime64('2016-05-01'), numpy.timedelta64(14100000000000,'ns'))\n",
    "    >>> print(_normalize_datetime('2016-05-01'))\n",
    "    (numpy.datetime64('2016-05-01'), numpy.timedelta64(0,'D'))\n",
    "    >>> x = pd.DataFrame({'x' : [np.datetime64('2015-01-01 05:00:00'), np.datetime64('2015-02-01 06:00:00')]})\n",
    "    >>> print(_normalize_datetime(x.x))\n",
    "    (array(['2015-01-01', '2015-02-01'], dtype='datetime64[D]'), array([18000000000000, 21600000000000], dtype='timedelta64[ns]'))\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",
    "    >>> print(_normalize_datetime(x.index))\n",
    "    (array(['2015-01-01', '2015-02-01'], dtype='datetime64[D]'), array([18000000000000, 21600000000000], dtype='timedelta64[ns]'))\n",
    "    '''\n",
    "    if isinstance(val, pd.Timestamp): \n",
    "        datetime = val.to_datetime64()\n",
    "    elif isinstance(val, pd.Series) or isinstance(val, pd.DatetimeIndex):\n",
    "        datetime = val.values\n",
    "    elif isinstance(val, np.ndarray) and np.issubdtype(val.dtype, np.datetime64): \n",
    "        datetime = val\n",
    "    else:\n",
    "        datetime = np.datetime64(val)\n",
    "        \n",
    "    date = datetime.astype('M8[D]')\n",
    "    time_delta = datetime - date\n",
    "    return date, time_delta\n",
    "\n",
    "\n",
    "def _normalize(start: Union[str, datetime.datetime, np.datetime64],\n",
    "               end: Union[str, datetime.datetime, np.datetime64],\n",
    "               include_first: bool,\n",
    "               include_last: bool) -> Tuple[np.datetime64, np.datetime64]:\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\n",
    "        end: end_date\n",
    "        include_first: whether to increment start date by 1\n",
    "        include_last: whether to increment end date by 1\n",
    "\n",
    "    Return:\n",
    "        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 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",
    "\n",
    "def read_holidays(calendar_name: str, dirname: str = None) -> np.ndarray:\n",
    "    '''\n",
    "    Reads a csv with a holidays column containing holidays (not including weekends)\n",
    "    '''\n",
    "    \n",
    "    curr_frame: FrameType = inspect.currentframe()  # type: ignore  # wants Optional[FrameType]\n",
    "    if dirname is None: dirname = os.path.dirname(os.path.abspath(inspect.getfile(curr_frame)))\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",
    "    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",
    "\n",
    "class Calendar:\n",
    "    \n",
    "    NYSE = \"nyse\"\n",
    "    EUREX = \"eurex\"\n",
    "    _calendars: MutableMapping[str, 'Calendar'] = {}\n",
    "    \n",
    "    def __init__(self, holidays: np.ndarray) -> None:\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: Union[np.datetime64, pd.Series, np.ndarray, str, datetime.datetime, datetime.date]) -> Union[bool, np.ndarray]:\n",
    " \n",
    "        '''\n",
    "        Returns whether the date is not a holiday or a weekend\n",
    "        \n",
    "        Args:\n",
    "            dates: date times to check\n",
    "        Return:\n",
    "            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, \n",
    "                         start: Union[np.datetime64, pd.Series, np.ndarray, str, datetime.datetime, datetime.date],\n",
    "                         end: Union[np.datetime64, pd.Series, np.ndarray, str, datetime.datetime, datetime.date],\n",
    "                         include_first: bool = False,\n",
    "                         include_last: bool = True) -> Union[int, np.ndarray]:\n",
    "        '''\n",
    "        Count the number of trading days between two date series including those two dates\n",
    "        \n",
    "        >>> eurex = Calendar.get_calendar(Calendar.EUREX)\n",
    "        >>> eurex.num_trading_days('2009-01-01', '2011-12-31')\n",
    "        772\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",
    "        iterable = isinstance(start, Iterable) and not isinstance(start, str)\n",
    "        s_tmp, e_tmp = _normalize(start, end, include_first, include_last)\n",
    "        # np.busday_count does not like nat dates\n",
    "        if iterable:\n",
    "            ret = np.full(len(s_tmp), np.nan)\n",
    "            mask = ~(np.isnat(s_tmp) | np.isnat(e_tmp))\n",
    "            count = np.busday_count(s_tmp[mask], e_tmp[mask], busdaycal=self.bus_day_cal)\n",
    "            ret[mask] = count\n",
    "            return ret\n",
    "        else:\n",
    "            if np.isnat(s_tmp) or np.isnat(s_tmp): return np.nan\n",
    "            count = np.busday_count(s_tmp, e_tmp, busdaycal=self.bus_day_cal)\n",
    "            return count\n",
    "        \n",
    "    def get_trading_days(self,\n",
    "                         start: Union[np.datetime64, pd.Series, np.ndarray, str, datetime.datetime, datetime.date],\n",
    "                         end: Union[np.datetime64, pd.Series, np.ndarray, str, datetime.datetime, datetime.date],\n",
    "                         include_first: bool = False,\n",
    "                         include_last: bool = True) -> Union[int, np.ndarray]:\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 third_friday_of_month(self, month: int, year: int, roll: str = 'backward') -> np.datetime64:\n",
    "        '''\n",
    "        >>> nyse = Calendar.get_calendar(Calendar.NYSE)\n",
    "        >>> nyse.third_friday_of_month(3, 2017)\n",
    "        numpy.datetime64('2017-03-17')\n",
    "        '''\n",
    "        # From https://stackoverflow.com/questions/18424467/python-third-friday-of-a-month\n",
    "        FRIDAY = 4\n",
    "        first_day_of_month = datetime.datetime(year, month, 1)\n",
    "        first_friday = first_day_of_month + datetime.timedelta(days=((FRIDAY - cal.monthrange(year, month)[0]) + 7) % 7)\n",
    "        # 4 is friday of week\n",
    "        third_friday_date = first_friday + datetime.timedelta(days=14)\n",
    "        third_friday = third_friday_date.date()\n",
    "        third_friday = self.add_trading_days(third_friday, 0, roll)\n",
    "        return third_friday\n",
    "    \n",
    "    def add_trading_days(self,\n",
    "                         start: Union[np.datetime64, pd.Series, np.ndarray, str, datetime.datetime, datetime.date],\n",
    "                         num_days: int, \n",
    "                         roll: str = 'raise') -> Union[np.datetime64, np.ndarray]:\n",
    "        '''\n",
    "        Adds trading days to a start date\n",
    "        \n",
    "        Args:\n",
    "            start: start datetimes(s)\n",
    "            num_days: number of trading days to add\n",
    "            roll: one of 'raise', 'nat', 'forward', 'following', 'backward', 'preceding', 'modifiedfollowing', 'modifiedpreceding' or 'allow'}\n",
    "                'allow' is a special case in which case, adding 1 day to a holiday will act as if it was not a holiday, and give you the next business day'\n",
    "                The rest of the values are the same as in the numpy busday_offset function\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",
    "        Return:\n",
    "            The datetime 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",
    "        >>> calendar.add_trading_days(np.datetime64('2019-02-17 15:25'), 1, roll = 'allow')\n",
    "        numpy.datetime64('2019-02-19T15:25')\n",
    "        >>> calendar.add_trading_days(np.datetime64('2019-02-17 15:25'), -1, roll = 'allow')\n",
    "        numpy.datetime64('2019-02-15T15:25')\n",
    "        '''\n",
    "        start_date, time_delta = _normalize_datetime(start)\n",
    "        if roll == 'allow':\n",
    "            # If today is a holiday, roll forward but subtract 1 day so\n",
    "            num_days = np.where(self.is_trading_day(start) | (num_days < 1), num_days, num_days - 1)\n",
    "            roll = 'forward'\n",
    "        out = np.busday_offset(start_date, num_days, roll=roll, busdaycal=self.bus_day_cal)\n",
    "        out = out + time_delta  # for some reason += does not work correctly here.\n",
    "        return out\n",
    "        \n",
    "    @staticmethod\n",
    "    def add_calendar(exchange_name: str, holidays: np.ndarray) -> None:\n",
    "        '''\n",
    "        Add a trading calendar to the class level calendars dict\n",
    "        \n",
    "        Args:\n",
    "            exchange_name: Name of the exchange.\n",
    "            holidays: holidays for this exchange, excluding weekends\n",
    "        '''\n",
    "        Calendar._calendars[exchange_name] = Calendar(holidays)\n",
    "        \n",
    "    @staticmethod\n",
    "    def get_calendar(exchange_name: str) -> 'Calendar':\n",
    "        '''\n",
    "        Get a calendar object for the given exchange:\n",
    "        \n",
    "        Args:\n",
    "            exchange_name: 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",
    "            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_name}')\n",
    "            holidays = read_holidays(exchange_name)\n",
    "            Calendar.add_calendar(exchange_name, holidays)\n",
    "        return Calendar._calendars[exchange_name]\n",
    "    \n",
    "\n",
    "def get_date_from_weekday(weekday: int, year: int, month: int, week: int) -> np.datetime64:\n",
    "    '''\n",
    "    Return the date that falls on a given weekday (Monday = 0) on a week, year and month\n",
    "    >>> get_date_from_weekday(1, 2019, 10, 4)\n",
    "    numpy.datetime64('2019-10-22')\n",
    "    '''\n",
    "    if week == -1:  # Last day of month\n",
    "        _, last_day = cal.monthrange(year, month)\n",
    "        return datetime.date(year, month, last_day)\n",
    "    first_day_of_month = datetime.date(year, month, 1)\n",
    "    date = first_day_of_month + rd.relativedelta(weeks=week - 1, weekday=weekday)\n",
    "    return np.datetime64(date)\n",
    "    \n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    import doctest\n",
    "    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
