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.

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

ENH: Improve general datetime functions #352

Merged
merged 15 commits into from
Oct 4, 2022
30 changes: 23 additions & 7 deletions pandas-stubs/core/indexes/datetimes.pyi
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
from datetime import (
date,
timedelta,
tzinfo,
)
from typing import (
Hashable,
Literal,
Sequence,
overload,
)

Expand All @@ -27,6 +28,7 @@ from pandas._typing import (
AnyArrayLike,
ArrayLike,
DatetimeLike,
IntervalClosedType,
)

from pandas.core.dtypes.dtypes import DatetimeTZDtype
Expand Down Expand Up @@ -90,24 +92,38 @@ def date_range(
start: str | DatetimeLike | None = ...,
end: str | DatetimeLike | None = ...,
periods: int | None = ...,
# TODO: Test timedelta and Timedelta, update pandas docs
# TODO: Add Timedelta to pandas docs
freq: str | timedelta | Timedelta | BaseOffset = ...,
tz: str | tzinfo = ...,
normalize: bool = ...,
name: Hashable | None = ...,
inclusive: Literal["left", "right"] | None = ...,
inclusive: IntervalClosedType = ...,
) -> DatetimeIndex: ...
@overload
def bdate_range(
start: str | DatetimeLike | None = ...,
end: str | DatetimeLike | None = ...,
periods: int | None = ...,
# TODO: Test timedelta and Timedelta, update pandas docs
freq: str | timedelta | Timedelta | BaseOffset = ...,
tz: str | tzinfo = ...,
normalize: bool = ...,
name: Hashable | None = ...,
weekmask: str | None = ...,
# TODO: Check if dt.date is allowed
holidays: list[str | DatetimeLike] | None = ...,
inclusive: Literal["left", "right"] | None = ...,
holidays: None = ...,
inclusive: IntervalClosedType = ...,
) -> DatetimeIndex: ...
@overload
def bdate_range(
start: str | DatetimeLike | None = ...,
end: str | DatetimeLike | None = ...,
periods: int | None = ...,
*,
# TODO: Add Timedelta to pandas docs
freq: str | timedelta | Timedelta | BaseOffset,
tz: str | tzinfo = ...,
normalize: bool = ...,
name: Hashable | None = ...,
weekmask: str | None = ...,
holidays: Sequence[str | DatetimeLike | date],
inclusive: IntervalClosedType = ...,
) -> DatetimeIndex: ...
4 changes: 3 additions & 1 deletion pandas-stubs/core/indexes/period.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@ from pandas.core.indexes.datetimelike import (
)
from pandas.core.indexes.numeric import Int64Index

from pandas._libs.tslibs import BaseOffset

class PeriodIndex(DatetimeIndexOpsMixin, Int64Index):
def __new__(
cls,
Expand Down Expand Up @@ -52,6 +54,6 @@ def period_range(
start: str | pd.Period | None = ...,
end: str | pd.Period | None = ...,
periods: int | None = ...,
freq: str | pd.DateOffset | None = ...,
freq: str | BaseOffset | None = ...,
Dr-Irv marked this conversation as resolved.
Show resolved Hide resolved
name: Hashable | None = ...,
) -> PeriodIndex: ...
51 changes: 51 additions & 0 deletions tests/test_timefuncs.py
Original file line number Diff line number Diff line change
Expand Up @@ -635,3 +635,54 @@ def test_to_timedelta_index() -> None:
check(
assert_type(pd.to_timedelta(arg6, "ms"), pd.TimedeltaIndex), pd.TimedeltaIndex
)


def test_bdate_range_holidays() -> None:
pd.bdate_range("2000-1-1", "2001-1-1", freq="C", holidays=["2000-12-15"])
pd.bdate_range(
"2000-1-1", "2001-1-1", freq="C", holidays=[pd.Timestamp(2000, 12, 15)]
)
pd.bdate_range(
"2000-1-1", "2001-1-1", freq="C", holidays=[np.datetime64("2000-12-15")]
)
pd.bdate_range(
"2000-1-1", "2001-1-1", freq="C", holidays=[dt.datetime(2000, 12, 15)]
)
pd.bdate_range(
"2000-1-1", "2001-1-1", freq="C", holidays=[dt.date(2000, 12, 15)], name=("a",)
)


def test_period_range() -> None:
check(
assert_type(
pd.period_range(pd.Period("2001Q1"), end=pd.Period("2010Q1")),
pd.PeriodIndex,
),
pd.PeriodIndex,
)
check(
assert_type(pd.period_range("2001Q1", end=pd.Period("2010Q1")), pd.PeriodIndex),
pd.PeriodIndex,
)
check(
assert_type(
pd.period_range("2001-01-01", end="2010-01-01", freq="Q"), pd.PeriodIndex
),
pd.PeriodIndex,
)
check(
assert_type(pd.period_range("2001Q1", periods=100, freq="Q"), pd.PeriodIndex),
pd.PeriodIndex,
)
check(
assert_type(pd.period_range("2001Q1", periods=100, freq=Day()), pd.PeriodIndex),
pd.PeriodIndex,
)
check(
assert_type(
pd.period_range("2001Q1", periods=100, freq=Day(), name=("A",)),
pd.PeriodIndex,
),
pd.PeriodIndex,
)