Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

TYP: Allow time manipulation functions to accept date and timedelta objects #20763

Merged
merged 1 commit into from Jan 8, 2022
Merged

TYP: Allow time manipulation functions to accept date and timedelta objects #20763

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
2 changes: 1 addition & 1 deletion numpy/__init__.pyi
View file
Open in desktop
Expand Up @@ -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_]: ...
Expand Down
71 changes: 36 additions & 35 deletions numpy/core/multiarray.pyi
View file
Open in desktop
Expand Up @@ -67,6 +67,7 @@ from numpy.typing import (
NDArray,
ArrayLike,
_SupportsArray,
_NestedSequence,
_FiniteNestedSequence,
_ArrayLikeBool_co,
_ArrayLikeUInt_co,
Expand Down Expand Up @@ -818,129 +819,129 @@ 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: ...

# `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 = ...,
Expand Down
12 changes: 12 additions & 0 deletions numpy/typing/tests/data/reveal/multiarray.pyi
View file
Open in desktop
@@ -1,3 +1,4 @@
import datetime as dt
from typing import Any, List, TypeVar
from pathlib import Path

Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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_]]

Expand Down