GH1327 Add overloead for date_range and timedelta_range

Author: loicdiridollou

pandas-stubs/core/groupby/generic.pyi

@@ -213,17 +213,17 @@ class SeriesGroupBy(GroupBy[Series[S2]], Generic[S2, ByT]):
 _TT = TypeVar("_TT", bound=Literal[True, False])
 
 # ty ignore needed because of https://github.com/astral-sh/ty/issues/157#issuecomment-3017337945
-class DFCallable1(Protocol[P]):  # ty: ignore[invalid-argument-type]
+class DFCallable1(Protocol[P]):
     def __call__(
         self, df: DataFrame, /, *args: P.args, **kwargs: P.kwargs
     ) -> Scalar | list | dict: ...
 
-class DFCallable2(Protocol[P]):  # ty: ignore[invalid-argument-type]
+class DFCallable2(Protocol[P]):
     def __call__(
         self, df: DataFrame, /, *args: P.args, **kwargs: P.kwargs
     ) -> DataFrame | Series: ...
 
-class DFCallable3(Protocol[P]):  # ty: ignore[invalid-argument-type]
+class DFCallable3(Protocol[P]):
     def __call__(self, df: Iterable, /, *args: P.args, **kwargs: P.kwargs) -> float: ...
 
 class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):

pandas-stubs/core/indexes/datetimes.pyi

@@ -102,11 +102,45 @@ class DatetimeIndex(
         self, periods: int = 1, freq: DateOffset | Timedelta | str | None = None
     ) -> Self: ...
 
+@overload
+def date_range(
+    end: str | DateAndDatetimeLike,
+    periods: int,
+    freq: str | timedelta | Timedelta | BaseOffset | None = None,
+    tz: TimeZones = None,
+    normalize: bool = False,
+    name: Hashable | None = None,
+    inclusive: IntervalClosedType = "both",
+    unit: TimeUnit | None = None,
+) -> DatetimeIndex: ...
+@overload
+def date_range(
+    start: str | DateAndDatetimeLike,
+    periods: int,
+    freq: str | timedelta | Timedelta | BaseOffset | None = None,
+    tz: TimeZones = None,
+    normalize: bool = False,
+    name: Hashable | None = None,
+    inclusive: IntervalClosedType = "both",
+    unit: TimeUnit | None = None,
+) -> DatetimeIndex: ...
+@overload
+def date_range(
+    start: str | DateAndDatetimeLike | None,
+    end: str | DateAndDatetimeLike | None,
+    freq: str | timedelta | Timedelta | BaseOffset | None = None,
+    tz: TimeZones = None,
+    normalize: bool = False,
+    name: Hashable | None = None,
+    inclusive: IntervalClosedType = "both",
+    unit: TimeUnit | None = None,
+) -> DatetimeIndex: ...
+@overload
 def date_range(
-    start: str | DateAndDatetimeLike | None = None,
-    end: str | DateAndDatetimeLike | None = None,
-    periods: int | None = None,
-    freq: str | timedelta | Timedelta | BaseOffset = "D",
+    start: str | DateAndDatetimeLike,
+    end: str | DateAndDatetimeLike,
+    periods: int,
+    freq: None = None,
     tz: TimeZones = None,
     normalize: bool = False,
     name: Hashable | None = None,

pandas-stubs/core/indexes/timedeltas.pyi

@@ -79,10 +79,51 @@ class TimedeltaIndex(
     def to_series(self, index=..., name: Hashable = ...) -> TimedeltaSeries: ...
     def shift(self, periods: int = 1, freq=...) -> Self: ...
 
+@overload
 def timedelta_range(
-    start: TimedeltaConvertibleTypes | None = None,
-    end: TimedeltaConvertibleTypes | None = None,
-    periods: int | None = None,
+    start: TimedeltaConvertibleTypes,
+    end: TimedeltaConvertibleTypes,
+    periods: int,
+    freq: None = None,
+    name: Hashable | None = None,
+    closed: Literal["left", "right"] | None = None,
+    *,
+    unit: None | str = ...,
+) -> TimedeltaIndex: ...
+@overload
+def timedelta_range(
+    start: TimedeltaConvertibleTypes,
+    periods: int,
+    freq: None = None,
+    name: Hashable | None = None,
+    closed: Literal["left", "right"] | None = None,
+    *,
+    unit: None | str = ...,
+) -> TimedeltaIndex: ...
+@overload
+def timedelta_range(
+    start: TimedeltaConvertibleTypes,
+    periods: int,
+    freq: str | DateOffset | Timedelta | dt.timedelta | None = None,
+    name: Hashable | None = None,
+    closed: Literal["left", "right"] | None = None,
+    *,
+    unit: None | str = ...,
+) -> TimedeltaIndex: ...
+@overload
+def timedelta_range(
+    start: TimedeltaConvertibleTypes,
+    end: TimedeltaConvertibleTypes,
+    freq: str | DateOffset | Timedelta | dt.timedelta | None = None,
+    name: Hashable | None = None,
+    closed: Literal["left", "right"] | None = None,
+    *,
+    unit: None | str = ...,
+) -> TimedeltaIndex: ...
+@overload
+def timedelta_range(
+    end: TimedeltaConvertibleTypes,
+    periods: int,
     freq: str | DateOffset | Timedelta | dt.timedelta | None = None,
     name: Hashable | None = None,
     closed: Literal["left", "right"] | None = None,

tests/test_styler.py

@@ -22,6 +22,7 @@
 from pandas._typing import Scalar
 
 from tests import (
+    PD_LTE_23,
     check,
     ensure_clean,
 )
@@ -141,7 +142,8 @@ def test_highlight_quantile() -> None:
 
 
 def test_loader() -> None:
-    check(assert_type(DF.style.loader, PackageLoader), PackageLoader)
+    if PD_LTE_23:
+        check(assert_type(DF.style.loader, PackageLoader), PackageLoader)
 
 
 def test_pipe() -> None:

tests/test_timefuncs.py

@@ -1563,6 +1563,13 @@ def test_timedelta_range() -> None:
         pd.TimedeltaIndex,
     )
 
+    if TYPE_CHECKING_INVALID_USAGE:
+        day1 = pd.Timedelta(1, unit="D")
+        day10 = pd.Timedelta(10, unit="D")
+        pd.timedelta_range(
+            day1, day10, periods=10, freq="D"  # type: ignore[call-overload]  # pyright: ignore[reportArgumentType]
+        )
+
 
 def test_dateoffset_freqstr() -> None:
     offset = DateOffset(minutes=10)