datelib.py

import datetime
from typing import Generator

import pytz
from dateutil import rrule


def iter_dates(
    start: datetime.datetime, end: datetime.datetime, day_delta=1
) -> Generator[datetime.datetime, None, None]:
    td = datetime.timedelta(days=day_delta)
    while start <= end:
        yield start
        start += td


def iter_eastern_10_am_one_year() -> Generator[datetime.datetime, None, None]:
    start = datetime.datetime(2021, 1, 1, 10, 0, tzinfo=pytz.timezone("US/Eastern"))
    end = datetime.datetime(2022, 1, 1, 10, 0, tzinfo=pytz.timezone("US/Eastern"))
    for date in iter_dates(start, end):
        yield date


def iter_market_dates_10_am_one_year() -> Generator[datetime.datetime, None, None]:
    for date in iter_eastern_10_am_one_year():
        if nyse.is_market(date):
            yield date


class NYSECalendar:
    def gen_holiday_ruleset_future(self, days):
        # gen a ruleset from today until days in future
        today = datetime.date.today()

        start = today
        end = today + datetime.timedelta(days)

        return self.gen_holiday_ruleset(start, end)

    def gen_holiday_ruleset_past(self, days):
        # gen a ruleset form today until n days in past
        today = datetime.date.today()

        start = today - datetime.timedelta(days)
        end = today

        return self.gen_holiday_ruleset(start, end)

    def gen_holiday_ruleset(self, start, end):
        """
        :param datetime.datetime start: the begin of the range range for ruleset
        :param datetime.datetime end: the end of the range for ruleset
        """
        rs = rrule.rruleset()

        # New Years Eve
        rs.rrule(
            rrule.rrule(
                rrule.YEARLY,
                dtstart=start,
                until=end,
                bymonth=12,
                bymonthday=31,
                byweekday=rrule.FR,
            )
        )

        # New Years Day
        rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start, until=end, bymonth=1, bymonthday=1))
        rs.rrule(
            rrule.rrule(
                rrule.YEARLY,
                dtstart=start,
                until=end,
                bymonth=1,
                bymonthday=2,
                byweekday=rrule.MO,
            )
        )

        # MLK
        rs.rrule(
            rrule.rrule(
                rrule.YEARLY, dtstart=start, until=end, bymonth=1, byweekday=rrule.MO(3)
            )
        )

        # Washington's Birtyday
        rs.rrule(
            rrule.rrule(
                rrule.YEARLY, dtstart=start, until=end, bymonth=2, byweekday=rrule.MO(3)
            )
        )

        # Good Friday
        rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start, until=end, byeaster=-2))

        # Memorial Day
        rs.rrule(
            rrule.rrule(
                rrule.YEARLY, dtstart=start, until=end, bymonth=5, byweekday=rrule.MO(-1)
            )
        )

        # July 4th
        rs.rrule(
            rrule.rrule(
                rrule.YEARLY,
                dtstart=start,
                until=end,
                bymonth=7,
                bymonthday=3,
                byweekday=rrule.FR,
            )
        )
        rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start, until=end, bymonth=7, bymonthday=4))
        rs.rrule(
            rrule.rrule(
                rrule.YEARLY,
                dtstart=start,
                until=end,
                bymonth=7,
                bymonthday=5,
                byweekday=rrule.MO,
            )
        )

        # Labor day
        rs.rrule(
            rrule.rrule(
                rrule.YEARLY, dtstart=start, until=end, bymonth=9, byweekday=rrule.MO(1)
            )
        )

        # Thanksgiving
        rs.rrule(
            rrule.rrule(
                rrule.YEARLY, dtstart=start, until=end, bymonth=11, byweekday=rrule.TH(4)
            )
        )

        # Christmas Eve
        rs.rrule(
            rrule.rrule(
                rrule.YEARLY,
                dtstart=start,
                until=end,
                bymonth=12,
                bymonthday=24,
                byweekday=rrule.FR,
            )
        )

        # Christmas
        rs.rrule(
            rrule.rrule(rrule.YEARLY, dtstart=start, until=end, bymonth=12, bymonthday=25)
        )
        rs.rrule(
            rrule.rrule(
                rrule.YEARLY,
                dtstart=start,
                until=end,
                bymonth=12,
                bymonthday=26,
                byweekday=rrule.MO,
            )
        )

        # Exclude potential holidays that fall on weekends
        rs.exrule(
            rrule.rrule(rrule.WEEKLY, dtstart=start, until=end, byweekday=(rrule.SA, rrule.SU))
        )

        return rs

    def gen_weekends_ruleset(self, start):
        rs = rrule.rruleset()
        rs.rrule(rrule.rrule(rrule.WEEKLY, dtstart=start, byweekday=(rrule.SA, rrule.SU)))

        return rs

    def gen_market_ruleset_past(self, days):
        today = datetime.date.today()

        start = today - datetime.timedelta(days)
        end = today

        return self.gen_market_ruleset(start, end)

    def gen_market_ruleset_future(self, days):
        today = datetime.date.today()

        start = today
        end = today + datetime.timedelta(days)

        return self.gen_market_ruleset(start, end)

    def gen_market_ruleset_relative(self, days, dt=None):
        # generate a ruleset relative today
        dt = dt or datetime.date.today()

        start = dt - datetime.timedelta(days)
        end = dt + datetime.timedelta(days)

        return self.gen_market_ruleset(start, end)

    def gen_market_ruleset(self, start, end):
        rs = rrule.rruleset()
        rs.rrule(rrule.rrule(rrule.DAILY, dtstart=start, until=end))

        weekends = self.gen_weekends_ruleset(start)  # noqa
        holidays = self.gen_holiday_ruleset(start, end)

        # update market ruleset to exclude weekends and the holiday schedule
        rs.exrule(rrule.rrule(rrule.WEEKLY, dtstart=start, byweekday=(rrule.SA, rrule.SU)))
        rs.exrule(holidays)

        return rs

    def is_tomorrow_market(self):
        today = datetime.datetime.now(tz=pytz.timezone("US/Eastern"))
        tomorrow = today + datetime.timedelta(1)
        rs = self.gen_market_ruleset_future(14)

        return self.contains(tomorrow, rs)

    def is_market(self, dt):
        # pass a datetime and return True/False if its a marketday
        dt_stripped = datetime.datetime.strptime(
            dt.strftime("%m-%d-%Y 0:00"), "%m-%d-%Y %H:%S"
        )
        rs = self.gen_market_ruleset_relative(14, dt=dt_stripped)

        return self.contains(dt, rs)

    def is_today_market(self):
        today = datetime.datetime.now(tz=pytz.timezone("US/Eastern"))
        rs = self.gen_market_ruleset_relative(14)

        return self.contains(today, rs)

    def is_yesterday_market(self):
        current_eastern_dt = datetime.datetime.now(tz=pytz.timezone("US/Eastern"))
        yesterday = current_eastern_dt - datetime.timedelta(1)
        rs = self.gen_market_ruleset_past(14)

        return self.contains(yesterday, rs)

    def is_now_market_hours(self):
        # market hours check with a 1m buffer at market close
        current_eastern_dt = datetime.datetime.now(tz=pytz.timezone("US/Eastern"))
        hour, minute = current_eastern_dt.strftime("%H:%M").split(":")
        hour, minute = int(hour), int(minute)
        if hour == 9:
            if minute < 30:
                return False
            if minute >= 30:  # 9:30 - 9:59
                return True

        if hour > 9 and hour < 15:  # 10:00 - 2:59
            return True

        if hour == 15:  # 3:00 - 3:58
            if minute < 58:
                return True

        return False

    def is_now_market_hours_fuzzy(self):
        # 6am-5pm market hours
        current_eastern_dt = datetime.datetime.now(tz=pytz.timezone("US/Eastern"))
        hour, minute = current_eastern_dt.strftime("%H:%M").split(":")
        hour, minute = int(hour), int(minute)
        if hour == 9:
            return True

        if hour > 9 and hour < 17:  # 10:00 - 4:59
            return True

        return False

    def get_market_time_simple(self):
        current_eastern_dt = datetime.datetime.now(tz=pytz.timezone("US/Eastern"))
        hour, minute = current_eastern_dt.strftime("%H:%M").split(":")
        hour, minute = int(hour), int(minute)

        return hour, minute

    def contains(self, dt, rs):
        # handle complex datetimes in rulesets
        dt_stripped = datetime.datetime.strptime(
            dt.strftime("%m-%d-%Y 0:00"), "%m-%d-%Y %H:%S"
        )

        return dt_stripped in rs


nyse = NYSECalendar()