Merge remote-tracking branch 'upstream/main' into typ-period

Author: MarcoGorelli

pandas-stubs/_typing.pyi

@@ -15,6 +15,7 @@ from re import Pattern
 import sys
 from typing import (
     Any,
+    Generic,
     Literal,
     Protocol,
     SupportsIndex,
@@ -37,6 +38,7 @@ from typing_extensions import (
     ParamSpec,
     TypeAlias,
     TypeVar,
+    override,
 )
 
 from pandas._libs.interval import Interval
@@ -818,6 +820,8 @@ np_ndarray_float: TypeAlias = npt.NDArray[np.floating]
 np_ndarray_complex: TypeAlias = npt.NDArray[np.complexfloating]
 np_ndarray_bool: TypeAlias = npt.NDArray[np.bool_]
 np_ndarray_str: TypeAlias = npt.NDArray[np.str_]
+np_ndarray_dt: TypeAlias = npt.NDArray[np.datetime64]
+np_ndarray_td: TypeAlias = npt.NDArray[np.timedelta64]
 
 # Define shape and generic type variables with defaults similar to numpy
 GenericT = TypeVar("GenericT", bound=np.generic, default=Any)
@@ -1055,4 +1059,14 @@ DictConvertible: TypeAlias = FulldatetimeDict | DataFrame
 # where it is the only acceptable type.
 Incomplete: TypeAlias = Any
 
+# differentiating between bool and int/float/complex
+# https://github.com/pandas-dev/pandas-stubs/pull/1312#pullrequestreview-3126128971
+class Just(Protocol, Generic[T]):
+    @property  # type: ignore[override]
+    @override
+    def __class__(self, /) -> type[T]: ...
+    @__class__.setter
+    @override
+    def __class__(self, t: type[T], /) -> None: ...
+
 __all__ = ["npt", "type_t"]

pandas-stubs/core/groupby/generic.pyi

@@ -264,6 +264,13 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):
         engine_kwargs: WindowingEngineKwargs = ...,
         **kwargs,
     ) -> DataFrame: ...
+    @overload
+    def aggregate(
+        self,
+        func: AggFuncTypeFrame | None = None,
+        /,
+        **kwargs,
+    ) -> DataFrame: ...
     agg = aggregate
     @overload
     def transform(

pandas-stubs/core/groupby/groupby.pyi

@@ -74,9 +74,9 @@ from pandas._typing import (
 from pandas.plotting import PlotAccessor
 
 _ResamplerGroupBy: TypeAlias = (
-    DatetimeIndexResamplerGroupby[NDFrameT]
-    | PeriodIndexResamplerGroupby[NDFrameT]
-    | TimedeltaIndexResamplerGroupby[NDFrameT]
+    DatetimeIndexResamplerGroupby[NDFrameT]  # ty: ignore[invalid-argument-type]
+    | PeriodIndexResamplerGroupby[NDFrameT]  # ty: ignore[invalid-argument-type]
+    | TimedeltaIndexResamplerGroupby[NDFrameT]  # ty: ignore[invalid-argument-type]
 )
 
 class GroupBy(BaseGroupBy[NDFrameT]):

pandas-stubs/core/indexes/datetimes.pyi

@@ -102,11 +102,46 @@ class DatetimeIndex(
         self, periods: int = 1, freq: DateOffset | Timedelta | str | None = None
     ) -> Self: ...
 
+@overload
+def date_range(
+    start: str | DateAndDatetimeLike,
+    end: str | DateAndDatetimeLike,
+    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,
+    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(
+    *,
+    end: str | DateAndDatetimeLike,
+    periods: int,
+    freq: str | timedelta | Timedelta | BaseOffset | None = None,
     tz: TimeZones = None,
     normalize: bool = False,
     name: Hashable | None = None,

pandas-stubs/core/indexes/timedeltas.pyi

@@ -79,13 +79,43 @@ 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,
+    *,
+    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,
+    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,
+    periods: int,
     *,
+    name: Hashable | None = None,
+    closed: Literal["left", "right"] | None = None,
     unit: None | str = ...,
 ) -> TimedeltaIndex: ...