MoFin/venv/lib/python3.12/site-packages/exchange_calendars/pandas_extensions/holiday.py

# https://github.com/pandas-dev/pandas/blob/master/pandas/tseries/holiday.py

import warnings

from pandas.tseries.holiday import (
    Holiday as PandasHoliday,
    AbstractHolidayCalendar as PandasAbstractHolidayCalendar,
)
from pandas.tseries.offsets import BaseOffset
from pandas import (
    DatetimeIndex,
    Series,
    Timestamp,
)
from pandas.errors import PerformanceWarning

from .offsets import _is_normalized


class Holiday(PandasHoliday):
    """
    This extension allows to have both offset and observance while the original
    Holiday does not. The order of application is offset => observance.
    """

    def __init__(
        self,
        name,
        year=None,
        month=None,
        day=None,
        offset=None,
        observance=None,
        start_date=None,
        end_date=None,
        days_of_week=None,
        tz=None,
    ):
        super().__init__(
            name,
            year,
            month,
            day,
            None,
            None,
            start_date,
            end_date,
            days_of_week,
        )
        self.offset = offset
        self.observance = observance
        self.tz = tz
        if self.start_date is not None:
            if self.tz is None and self.start_date.tz is not None:
                self.tz = start_date.tz
            if self.start_date.tz is None and self.tz is not None:
                self.start_date = self.start_date.tz_localize(self.tz)
            assert self.tz == self.start_date.tz
        if self.end_date is not None:
            if self.tz is None and self.end_date.tz is not None:
                self.tz = end_date.tz
            if self.end_date.tz is None and self.tz is not None:
                self.end_date = self.end_date.tz_localize(self.tz)
            assert self.tz == self.end_date.tz
        if self.start_date is not None and self.end_date is not None:
            assert self.start_date.tz == self.end_date.tz

    def __repr__(self) -> str:
        info = ""
        if self.year is not None:
            info += f"year={self.year}, "
        info += f"month={self.month}, day={self.day}"

        if self.offset is not None:
            info += f", offset={self.offset}"

        if self.observance is not None:
            info += f", observance={self.observance}"

        return f"{self.__class__.__name__}: {self.name} ({info})"

    def dates(self, start_date, end_date, return_name=False):
        start_date = Timestamp(start_date)
        end_date = Timestamp(end_date)
        assert start_date.tz == self.tz
        assert _is_normalized(start_date)
        assert end_date.tz == self.tz
        assert _is_normalized(end_date)
        return super().dates(start_date, end_date, return_name=return_name)

    def _apply_offset(self, dates):
        if self.offset is not None:
            if not isinstance(self.offset, list):
                offsets = [self.offset]
            else:
                offsets = self.offset
            for offset in offsets:
                # If we are adding a non-vectorized value
                # ignore the PerformanceWarnings:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore", PerformanceWarning)
                    dates += offset
        return dates

    def _apply_observance(self, dates):
        if self.observance is not None:
            if not isinstance(self.observance, list):
                observances = [self.observance]
            else:
                observances = self.observance
            for observance in observances:
                if isinstance(observance, BaseOffset):
                    offset = observance

                    def f(d):
                        return d + offset  # noqa: B023

                    observance = f  # noqa: PLW2901
                dates = dates.map(observance)
        return dates

    def _apply_rule(self, dates, apply_offset=True, apply_observance=True):
        if apply_offset:
            dates = self._apply_offset(dates)
        if apply_observance:
            dates = self._apply_observance(dates)
        return dates


def combine(pre_holidays):
    combined = Series(index=DatetimeIndex([]), dtype=object)
    for holidays in pre_holidays:
        combined = combined.combine_first(holidays)
    return combined


class AbstractHolidayCalendar(PandasAbstractHolidayCalendar):
    """
    This extension allows to have overlaps between calculated holidays while
    the original AbstractHolidayCalendar does not.
    """

    def holidays(self, start=None, end=None, return_name=False):
        """
        Returns a curve with holidays between start_date and end_date
        Parameters
        ----------
        start : starting date, datetime-like, optional
        end : ending date, datetime-like, optional
        return_name : bool, optional
            If True, return a series that has dates and holiday names.
            False will only return a DatetimeIndex of dates.
        Returns
        -------
            DatetimeIndex of holidays
        """
        if self.rules is None:
            raise Exception(
                f"Holiday Calendar {self.name} does not have any rules specified"
            )

        if start is None:
            start = AbstractHolidayCalendar.start_date

        if end is None:
            end = AbstractHolidayCalendar.end_date

        start = Timestamp(start)
        end = Timestamp(end)

        # If we don't have a cache or the dates are outside the prior cache, we
        # get them again
        if self._cache is None or start < self._cache[0] or end > self._cache[1]:
            pre_holidays = [
                rule.dates(start, end, return_name=True) for rule in self.rules
            ]
            if pre_holidays:
                # This line's behavior is changed to use custom combine()
                # function instead of the original concat() function
                holidays = combine(pre_holidays)
            else:
                holidays = Series(index=DatetimeIndex([]), dtype=object)

            self._cache = (start, end, holidays.sort_index())

        holidays = self._cache[2]
        holidays = holidays[start:end]

        if return_name:
            return holidays
        return holidays.index