Trading days offset

trading_day.py

from typing import Union

import pandas as pd
from trading_calendars import TradingCalendar, get_calendar


class TradingDay:
    """Offset for trading days

    The class is useful to find previous and next trading days from given date.

    Examples
    --------

    >>> from trading_calendars import get_calendar
    >>> cal = get_calendar('XMOS')  # Moscow Exchange

    >>> '2018-01-09' - TradingDay(cal, days=2)
    Timestamp('2018-01-04')

    >>> '2018-01-09' + TradingDay(cal, days=-2)
    Timestamp('2018-01-04')

    >>> '2018-01-04' + TradingDay(cal, days=2)
    Timestamp('2018-01-09')

    Parameters
    ----------
    calendar : Union[str, TradingCalendar]
        The trading calendar for the exchange (see trading-calendars package)

    days : int
        The number of trading days (offset from the given date). It can be positive or negative integer.

    """

    def __init__(self, calendar: Union[str, TradingCalendar], days: int):
        if isinstance(calendar, str):
            calendar = get_calendar(calendar)

        self._calendar = calendar
        self._days = days

    @property
    def calendar(self) -> TradingCalendar:
        return self._calendar

    @property
    def days(self) -> int:
        return self._days

    def __repr__(self) -> str:
        return f"{type(self).__name__}(calendar='{self.calendar.name}', days={self.days})"

    def __neg__(self) -> 'TradingDay':
        return TradingDay(self._calendar, days=-self._days)

    def __rsub__(self, other: object) -> pd.Timestamp:
        try:
            d = pd.Timestamp(other)
        except (TypeError, ValueError):
            return NotImplemented

        if self._days == 0:
            return d
        elif self._days < 0:
            return self._add(d)
        else:
            return self._sub(d)

    def __radd__(self, other: object) -> pd.Timestamp:
        try:
            d = pd.Timestamp(other)
        except (TypeError, ValueError):
            return NotImplemented

        if self._days == 0:
            return d
        elif self._days < 0:
            return self._sub(d)
        else:
            return self._add(d)

    def _sub(self, d):
        # Fix time to correct find the previous trading day
        d = pd.Timestamp(year=d.year, month=d.month, day=d.day, hour=0, minute=0, tz=d.tz)
        tz = d.tz

        for _ in range(abs(self._days)):
            d = pd.Timestamp(self._calendar.previous_open(d))

        return pd.Timestamp(d.date(), tz=tz)

    def _add(self, d):
        # Fix time to correct find the next trading day
        d = pd.Timestamp(year=d.year, month=d.month, day=d.day, hour=23, minute=59, tz=d.tz)
        tz = d.tz

        for _ in range(abs(self._days)):
            d = pd.Timestamp(self._calendar.next_open(d))

        return pd.Timestamp(d.date(), tz=tz)