Last active
January 24, 2021 15:39
-
-
Save espdev/135055b78cdad709025061b920578c5b to your computer and use it in GitHub Desktop.
Trading days offset
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment