diff --git a/pandas-stubs/_libs/tslibs/offsets.pyi b/pandas-stubs/_libs/tslibs/offsets.pyi index 852a24591..1b86a3fed 100644 --- a/pandas-stubs/_libs/tslibs/offsets.pyi +++ b/pandas-stubs/_libs/tslibs/offsets.pyi @@ -1,4 +1,5 @@ from datetime import ( + date, datetime, timedelta, ) @@ -15,10 +16,12 @@ from pandas.core.indexes.datetimes import DatetimeIndex from pandas._typing import npt +from pandas.tseries.holiday import AbstractHolidayCalendar + from .timedeltas import Timedelta _BaseOffsetT = TypeVar("_BaseOffsetT", bound=BaseOffset) -_DatetimeT = TypeVar("_DatetimeT", bound=datetime) +_DatetimeT = TypeVar("_DatetimeT", bound=date) _TimedeltaT = TypeVar("_TimedeltaT", bound=timedelta) prefix_mapping: dict[str, type] @@ -204,8 +207,8 @@ class CustomBusinessDay(BusinessDay): self, n: int = ..., normalize: bool = ..., - offset: timedelta = ..., - weekmask: str = ..., + holidays: list = ..., + calendar: AbstractHolidayCalendar | np.busdaycalendar = ..., ): ... class CustomBusinessHour(BusinessHour): diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 04f051bd8..73617ff4f 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -47,6 +47,8 @@ PandasScalar = Union[bytes, datetime.date, datetime.datetime, datetime.timedelta # Scalar = Union[PythonScalar, PandasScalar] IntStrT = TypeVar("IntStrT", int, str) +DatetimeLike = Union[datetime.date, datetime.datetime, np.datetime64, Timestamp] + # dtypes NpDtype = Union[str, np.dtype[np.generic], type[Union[str, complex, bool, object]]] Dtype = Union[ExtensionDtype, NpDtype] diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi index 32f9d442d..b578eeb77 100644 --- a/pandas-stubs/core/indexes/datetimes.pyi +++ b/pandas-stubs/core/indexes/datetimes.pyi @@ -19,11 +19,15 @@ from pandas.core.series import ( from pandas._typing import ( AnyArrayLike, ArrayLike, + DatetimeLike, + IntervalClosedType, np_ndarray_bool, ) from pandas.core.dtypes.dtypes import DatetimeTZDtype +from pandas.tseries.offsets import BaseOffset + class DatetimeIndex(DatetimeTimedeltaMixin, DatetimeIndexProperties): tz: tzinfo | None def __init__( @@ -76,25 +80,26 @@ class DatetimeIndex(DatetimeTimedeltaMixin, DatetimeIndexProperties): def dtype(self) -> np.dtype | DatetimeTZDtype: ... def date_range( - start=..., - end=..., - periods=..., - freq=..., - tz=..., - normalize=..., - name=..., - closed=..., + start: str | DatetimeLike | None = ..., + end: str | DatetimeLike | None = ..., + periods: int | None = ..., + freq: str | BaseOffset = ..., + tz: str | tzinfo = ..., + normalize: bool = ..., + name: str | None = ..., + inclusive: IntervalClosedType = ..., **kwargs, ) -> DatetimeIndex: ... def bdate_range( - start=..., - end=..., - periods=..., - freq: str = ..., - tz=..., + start: str | DatetimeLike | None = ..., + end: str | DatetimeLike | None = ..., + periods: int | None = ..., + freq: str | BaseOffset = ..., + tz: str | tzinfo = ..., normalize: bool = ..., - name=..., - weekmask=..., - holidays=..., - closed=..., + name: str | None = ..., + weekmask: str | None = ..., + holidays: list | None = ..., + inclusive: IntervalClosedType = ..., + **kwargs, ) -> DatetimeIndex: ... diff --git a/pandas-stubs/tseries/holiday.pyi b/pandas-stubs/tseries/holiday.pyi new file mode 100644 index 000000000..47c088e05 --- /dev/null +++ b/pandas-stubs/tseries/holiday.pyi @@ -0,0 +1,3 @@ +class Holiday: ... +class AbstractHolidayCalendar: ... +class USFederalHolidayCalendar(AbstractHolidayCalendar): ... diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 3e082e7da..c2ba18200 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -23,6 +23,13 @@ check, ) +from pandas.tseries.holiday import USFederalHolidayCalendar +from pandas.tseries.offsets import ( + BusinessDay, + CustomBusinessDay, + Day, +) + if TYPE_CHECKING: from pandas.core.series import ( PeriodSeries, @@ -403,3 +410,33 @@ def test_datetimeindex_accessors() -> None: check(assert_type(i0.month_name(), pd.Index), pd.Index, str) check(assert_type(i0.day_name(), pd.Index), pd.Index, str) check(assert_type(i0.is_normalized, bool), bool) + + +def test_some_offsets() -> None: + # GH 222 + + check( + assert_type( + CustomBusinessDay(calendar=USFederalHolidayCalendar()), CustomBusinessDay + ), + CustomBusinessDay, + ) + # GH 223 + check( + assert_type(pd.date_range("1/1/2022", "2/1/2022", freq="1D"), pd.DatetimeIndex), + pd.DatetimeIndex, + ) + check( + assert_type( + pd.date_range("1/1/2022", "2/1/2022", freq=Day()), pd.DatetimeIndex + ), + pd.DatetimeIndex, + ) + check( + assert_type( + pd.bdate_range("1/1/2022", "2/1/2022", freq=BusinessDay()), pd.DatetimeIndex + ), + pd.DatetimeIndex, + ) + # GH 224 + check(assert_type(dt.date.today() - Day(), dt.date), dt.date)