summaryrefslogtreecommitdiff
path: root/numpy
diff options
context:
space:
mode:
Diffstat (limited to 'numpy')
-rw-r--r--numpy/__init__.pyi2
-rw-r--r--numpy/core/multiarray.pyi71
-rw-r--r--numpy/typing/tests/data/reveal/multiarray.pyi12
3 files changed, 49 insertions, 36 deletions
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index 0facafaa8..b60d47790 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -3367,7 +3367,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 e06220dc3..423aed85e 100644
--- a/numpy/core/multiarray.pyi
+++ b/numpy/core/multiarray.pyi
@@ -61,6 +61,7 @@ from numpy.typing import (
NDArray,
ArrayLike,
_SupportsArray,
+ _NestedSequence,
_FiniteNestedSequence,
_ArrayLikeBool_co,
_ArrayLikeUInt_co,
@@ -812,28 +813,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: ...
@@ -841,100 +842,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 cbe30c520..f8937c379 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, 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_]]
View Lorry Controller Status