diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 9f3e42a69957..9cd67c272f5c 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -3374,7 +3374,7 @@ class busdaycalendar: def __new__( cls, weekmask: ArrayLike = ..., - holidays: ArrayLike = ..., + holidays: ArrayLike | dt.date | _NestedSequence[dt.date] = ..., ) -> busdaycalendar: ... @property def weekmask(self) -> NDArray[bool_]: ... diff --git a/numpy/core/multiarray.pyi b/numpy/core/multiarray.pyi index a9f68e1815cf..3c1c37108fdf 100644 --- a/numpy/core/multiarray.pyi +++ b/numpy/core/multiarray.pyi @@ -67,6 +67,7 @@ from numpy.typing import ( NDArray, ArrayLike, _SupportsArray, + _NestedSequence, _FiniteNestedSequence, _ArrayLikeBool_co, _ArrayLikeUInt_co, @@ -818,28 +819,28 @@ def datetime_data( @overload def busday_count( # type: ignore[misc] - begindates: _ScalarLike_co, - enddates: _ScalarLike_co, + begindates: _ScalarLike_co | dt.date, + enddates: _ScalarLike_co | dt.date, weekmask: ArrayLike = ..., - holidays: None | ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., busdaycal: None | busdaycalendar = ..., out: None = ..., ) -> int_: ... @overload def busday_count( # type: ignore[misc] - begindates: ArrayLike, - enddates: ArrayLike, + begindates: ArrayLike | dt.date | _NestedSequence[dt.date], + enddates: ArrayLike | dt.date | _NestedSequence[dt.date], weekmask: ArrayLike = ..., - holidays: None | ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., busdaycal: None | busdaycalendar = ..., out: None = ..., ) -> NDArray[int_]: ... @overload def busday_count( - begindates: ArrayLike, - enddates: ArrayLike, + begindates: ArrayLike | dt.date | _NestedSequence[dt.date], + enddates: ArrayLike | dt.date | _NestedSequence[dt.date], weekmask: ArrayLike = ..., - holidays: None | ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., busdaycal: None | busdaycalendar = ..., out: _ArrayType = ..., ) -> _ArrayType: ... @@ -847,100 +848,100 @@ def busday_count( # `roll="raise"` is (more or less?) equivalent to `casting="safe"` @overload def busday_offset( # type: ignore[misc] - dates: datetime64, - offsets: _TD64Like_co, + dates: datetime64 | dt.date, + offsets: _TD64Like_co | dt.timedelta, roll: L["raise"] = ..., weekmask: ArrayLike = ..., - holidays: None | ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., busdaycal: None | busdaycalendar = ..., out: None = ..., ) -> datetime64: ... @overload def busday_offset( # type: ignore[misc] - dates: _ArrayLike[datetime64], - offsets: _ArrayLikeTD64_co, + dates: _ArrayLike[datetime64] | dt.date | _NestedSequence[dt.date], + offsets: _ArrayLikeTD64_co | dt.timedelta | _NestedSequence[dt.timedelta], roll: L["raise"] = ..., weekmask: ArrayLike = ..., - holidays: None | ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., busdaycal: None | busdaycalendar = ..., out: None = ..., ) -> NDArray[datetime64]: ... @overload def busday_offset( # type: ignore[misc] - dates: _ArrayLike[datetime64], - offsets: _ArrayLike[timedelta64], + dates: _ArrayLike[datetime64] | dt.date | _NestedSequence[dt.date], + offsets: _ArrayLikeTD64_co | dt.timedelta | _NestedSequence[dt.timedelta], roll: L["raise"] = ..., weekmask: ArrayLike = ..., - holidays: None | ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., busdaycal: None | busdaycalendar = ..., out: _ArrayType = ..., ) -> _ArrayType: ... @overload def busday_offset( # type: ignore[misc] - dates: _ScalarLike_co, - offsets: _ScalarLike_co, + dates: _ScalarLike_co | dt.date, + offsets: _ScalarLike_co | dt.timedelta, roll: _RollKind, weekmask: ArrayLike = ..., - holidays: None | ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., busdaycal: None | busdaycalendar = ..., out: None = ..., ) -> datetime64: ... @overload def busday_offset( # type: ignore[misc] - dates: ArrayLike, - offsets: ArrayLike, + dates: ArrayLike | dt.date | _NestedSequence[dt.date], + offsets: ArrayLike | dt.timedelta | _NestedSequence[dt.timedelta], roll: _RollKind, weekmask: ArrayLike = ..., - holidays: None | ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., busdaycal: None | busdaycalendar = ..., out: None = ..., ) -> NDArray[datetime64]: ... @overload def busday_offset( - dates: ArrayLike, - offsets: ArrayLike, + dates: ArrayLike | dt.date | _NestedSequence[dt.date], + offsets: ArrayLike | dt.timedelta | _NestedSequence[dt.timedelta], roll: _RollKind, weekmask: ArrayLike = ..., - holidays: None | ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., busdaycal: None | busdaycalendar = ..., out: _ArrayType = ..., ) -> _ArrayType: ... @overload def is_busday( # type: ignore[misc] - dates: _ScalarLike_co, + dates: _ScalarLike_co | dt.date, weekmask: ArrayLike = ..., - holidays: None | ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., busdaycal: None | busdaycalendar = ..., out: None = ..., ) -> bool_: ... @overload def is_busday( # type: ignore[misc] - dates: ArrayLike, + dates: ArrayLike | _NestedSequence[dt.date], weekmask: ArrayLike = ..., - holidays: None | ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., busdaycal: None | busdaycalendar = ..., out: None = ..., ) -> NDArray[bool_]: ... @overload def is_busday( - dates: ArrayLike, + dates: ArrayLike | _NestedSequence[dt.date], weekmask: ArrayLike = ..., - holidays: None | ArrayLike = ..., + holidays: None | ArrayLike | dt.date | _NestedSequence[dt.date] = ..., busdaycal: None | busdaycalendar = ..., out: _ArrayType = ..., ) -> _ArrayType: ... @overload def datetime_as_string( # type: ignore[misc] - arr: datetime64, + arr: datetime64 | dt.date, unit: None | L["auto"] | _UnitKind = ..., timezone: L["naive", "UTC", "local"] | dt.tzinfo = ..., casting: _CastingKind = ..., ) -> str_: ... @overload def datetime_as_string( - arr: _ArrayLikeDT64_co, + arr: _ArrayLikeDT64_co | _NestedSequence[dt.date], unit: None | L["auto"] | _UnitKind = ..., timezone: L["naive", "UTC", "local"] | dt.tzinfo = ..., casting: _CastingKind = ..., diff --git a/numpy/typing/tests/data/reveal/multiarray.pyi b/numpy/typing/tests/data/reveal/multiarray.pyi index 0e91a7afdc52..56706d8d6f91 100644 --- a/numpy/typing/tests/data/reveal/multiarray.pyi +++ b/numpy/typing/tests/data/reveal/multiarray.pyi @@ -1,3 +1,4 @@ +import datetime as dt from typing import Any, List, TypeVar from pathlib import Path @@ -27,6 +28,10 @@ b_i8_f8_f8 = np.broadcast(AR_i8, AR_f8, AR_f8) nditer_obj: np.nditer +date_scalar: dt.date +date_seq: list[dt.date] +timedelta_seq: list[dt.timedelta] + def func(a: int) -> bool: ... reveal_type(next(b_f8)) # E: tuple[Any] @@ -108,19 +113,26 @@ reveal_type(np.datetime_data(np.dtype(np.timedelta64))) # E: Tuple[builtins.str reveal_type(np.busday_count("2011-01", "2011-02")) # E: {int_} reveal_type(np.busday_count(["2011-01"], "2011-02")) # E: ndarray[Any, dtype[{int_}]] +reveal_type(np.busday_count(["2011-01"], date_scalar)) # E: ndarray[Any, dtype[{int_}]] reveal_type(np.busday_offset(M, m)) # E: datetime64 +reveal_type(np.busday_offset(date_scalar, m)) # E: datetime64 reveal_type(np.busday_offset(M, 5)) # E: datetime64 reveal_type(np.busday_offset(AR_M, m)) # E: ndarray[Any, dtype[datetime64]] +reveal_type(np.busday_offset(M, timedelta_seq)) # E: ndarray[Any, dtype[datetime64]] reveal_type(np.busday_offset("2011-01", "2011-02", roll="forward")) # E: datetime64 reveal_type(np.busday_offset(["2011-01"], "2011-02", roll="forward")) # E: ndarray[Any, dtype[datetime64]] reveal_type(np.is_busday("2012")) # E: bool_ +reveal_type(np.is_busday(date_scalar)) # E: bool_ reveal_type(np.is_busday(["2012"])) # E: ndarray[Any, dtype[bool_]] reveal_type(np.datetime_as_string(M)) # E: str_ reveal_type(np.datetime_as_string(AR_M)) # E: ndarray[Any, dtype[str_]] +reveal_type(np.busdaycalendar(holidays=date_seq)) # E: busdaycalendar +reveal_type(np.busdaycalendar(holidays=[M])) # E: busdaycalendar + reveal_type(np.compare_chararrays("a", "b", "!=", rstrip=False)) # E: ndarray[Any, dtype[bool_]] reveal_type(np.compare_chararrays(b"a", b"a", "==", True)) # E: ndarray[Any, dtype[bool_]]