fix: #718 mean

Author: cmp0xff

pandas-stubs/core/series.pyi

@@ -25,6 +25,7 @@ from typing import (
     Generic,
     Literal,
     NoReturn,
+    TypeVar,
     overload,
 )
 
@@ -181,6 +182,13 @@ from pandas.core.dtypes.dtypes import CategoricalDtype
 
 from pandas.plotting import PlotAccessor
 
+_scalar_timedelta: TypeAlias = timedelta | np.timedelta64 | BaseOffset | Timedelta
+_vector_timedelta: TypeAlias = (
+    Sequence[timedelta] | Sequence[Timedelta] | Series[Timedelta] | TimedeltaIndex
+)
+_all_timedelta: TypeAlias = _scalar_timedelta | _vector_timedelta
+_stamp_and_delta = TypeVar("_stamp_and_delta", bound=Timestamp | Timedelta)
+
 class _iLocIndexerSeries(_iLocIndexer, Generic[S1]):
     # get item
     @overload
@@ -1591,15 +1599,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
     ) -> Series[int]: ...
     @overload
     def __add__(
-        self: Series[Timestamp],
-        other: (
-            Timedelta
-            | Series[Timedelta]
-            | TimedeltaIndex
-            | np.timedelta64
-            | timedelta
-            | BaseOffset
-        ),
+        self: Series[Timestamp], other: _all_timedelta
     ) -> Series[Timestamp]: ...
     @overload
     def __add__(self: Series[Timedelta], other: Period) -> PeriodSeries: ...
@@ -1686,7 +1686,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
     def __or__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ...
     @overload
     def __radd__(
-        self: Series[Timestamp], other: Series[Timedelta] | np.timedelta64 | timedelta
+        self: Series[Timestamp], other: _all_timedelta
     ) -> Series[Timestamp]: ...
     @overload
     def __radd__(
@@ -1707,14 +1707,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
     def __rdivmod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ...  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
     @overload
     def __rfloordiv__(
-        self: Series[Timedelta],
-        other: (
-            timedelta
-            | Series[Timedelta]
-            | np.timedelta64
-            | TimedeltaIndex
-            | Sequence[timedelta]
-        ),
+        self: Series[Timedelta], other: _all_timedelta
     ) -> Series[int]: ...
     @overload
     def __rfloordiv__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ...
@@ -1737,16 +1730,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
     def __rsub__(self, other: num | _ListLike | Series[S1]) -> Series: ...
     @overload
     def __rtruediv__(
-        self: Series[Timedelta],
-        other: (
-            timedelta
-            | Timedelta
-            | Sequence[Timedelta]
-            | Series[Timedelta]
-            | np.timedelta64
-            | TimedeltaIndex
-            | Sequence[timedelta]
-        ),
+        self: Series[Timedelta], other: _all_timedelta
     ) -> Series[float]: ...
     @overload
     def __rtruediv__(self, other: num | _ListLike | Series[S1] | Path) -> Series: ...
@@ -1764,15 +1748,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
     ) -> Series[int]: ...
     @overload
     def __sub__(
-        self: Series[Timestamp],
-        other: (
-            Timedelta
-            | Series[Timedelta]
-            | TimedeltaIndex
-            | np.timedelta64
-            | timedelta
-            | BaseOffset
-        ),
+        self: Series[Timestamp], other: _all_timedelta
     ) -> Series[Timestamp]: ...
     @overload
     def __sub__(
@@ -1794,16 +1770,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
     ) -> Series[Timestamp]: ...
     @overload
     def __truediv__(
-        self: Series[Timedelta],
-        other: (
-            timedelta
-            | Timedelta
-            | Sequence[Timedelta]
-            | Series[Timedelta]
-            | np.timedelta64
-            | TimedeltaIndex
-            | Sequence[timedelta]
-        ),
+        self: Series[Timedelta], other: _all_timedelta
     ) -> Series[float]: ...
     @overload
     def __truediv__(
@@ -1996,25 +1963,16 @@ class Series(IndexOpsMixin[S1], NDFrame):
         numeric_only: _bool = ...,
         **kwargs: Any,
     ) -> S1: ...
-    # @overload
-    # def mean(
-    #     self: Series[Timestamp],
-    #     axis: AxisIndex | None = ...,
-    #     skipna: _bool = ...,
-    #     level: None = ...,
-    #     numeric_only: _bool = ...,
-    #     **kwargs: Any,
-    # ) -> Timestamp: ...
-    # @overload
-    # def mean(
-    #     self: Series[Timedelta],
-    #     axis: AxisIndex | None = ...,
-    #     skipna: _bool = ...,
-    #     level: None = ...,
-    #     numeric_only: _bool = ...,
-    #     **kwargs: Any,
-    # ) -> Timedelta: ...
-    # @overload
+    @overload
+    def mean(
+        self: Series[_stamp_and_delta],
+        axis: AxisIndex | None = ...,
+        skipna: _bool = ...,
+        level: None = ...,
+        numeric_only: _bool = ...,
+        **kwargs: Any,
+    ) -> _stamp_and_delta: ...
+    @overload
     def mean(
         self,
         axis: AxisIndex | None = ...,

tests/test_timefuncs.py

@@ -400,10 +400,7 @@ def test_series_dt_accessors() -> None:
         upper="2.3.99",
     ):
         check(
-            assert_type(
-                s0.dt.to_pydatetime(),
-                np.ndarray,
-            ),
+            assert_type(s0.dt.to_pydatetime(), np.ndarray),
             np.ndarray if PD_LTE_23 else pd.Series,
             dt.datetime,
         )
@@ -598,13 +595,7 @@ def test_series_dt_accessors() -> None:
         "The behavior of TimedeltaProperties.to_pytimedelta is deprecated",
         lower="2.3.99",
     ):
-        check(
-            assert_type(
-                s2.dt.to_pytimedelta(),
-                np.ndarray,
-            ),
-            np.ndarray,
-        )
+        check(assert_type(s2.dt.to_pytimedelta(), np.ndarray), np.ndarray)
     check(assert_type(s2.dt.total_seconds(), "pd.Series[float]"), pd.Series, float)
     check(assert_type(s2.dt.unit, TimeUnit), str)
     check(
@@ -1497,7 +1488,7 @@ def test_timedeltaseries_add_timestampseries() -> None:
 
 def test_mean_median_std() -> None:
     s = pd.Series([pd.Timedelta("1 ns"), pd.Timedelta("2 ns"), pd.Timedelta("3 ns")])
-    # check(assert_type(s.mean(), pd.Timedelta), pd.Timedelta)
+    check(assert_type(s.mean(), pd.Timedelta), pd.Timedelta)
     check(assert_type(s.median(), pd.Timedelta), pd.Timedelta)
     check(assert_type(s.std(), pd.Timedelta), pd.Timedelta)
 
@@ -1508,7 +1499,7 @@ def test_mean_median_std() -> None:
             pd.Timestamp("2021-01-03"),
         ]
     )
-    # check(assert_type(s2.mean(), pd.Timestamp), pd.Timestamp)
+    check(assert_type(s2.mean(), pd.Timestamp), pd.Timestamp)
     check(assert_type(s2.median(), pd.Timestamp), pd.Timestamp)
     check(assert_type(s2.std(), pd.Timedelta), pd.Timedelta)