feat: comments

Author: cmp0xff

pandas-stubs/core/series.pyi

@@ -27,6 +27,7 @@ from typing import (
     Literal,
     NoReturn,
     Protocol,
+    TypeVar,
     final,
     overload,
     type_check_only,
@@ -210,10 +211,16 @@ from pandas.core.dtypes.dtypes import CategoricalDtype
 
 from pandas.plotting import PlotAccessor
 
+_T_INTERVAL_NP = TypeVar("_T_INTERVAL_NP", bound=np.bytes_ | np.str_)
+
 @type_check_only
 class _SupportsAdd(Protocol[_T_co]):
     def __add__(self, value: Self, /) -> _T_co: ...
 
+@type_check_only
+class SupportsSelfSub(Protocol[_T_co]):
+    def __sub__(self, x: Self, /) -> _T_co: ...
+
 @type_check_only
 class _SupportsMul(Protocol[_T_co]):
     def __mul__(self, value: Self, /) -> _T_co: ...
@@ -832,22 +839,19 @@ class Series(IndexOpsMixin[S1], NDFrame):
         self, other: Series[S1], min_periods: int | None = None, ddof: int = 1
     ) -> float: ...
     @overload
-    def diff(self: Series[_bool], periods: int = ...) -> Series[type[object]]: ...  # type: ignore[overload-overlap]
-    @overload
-    def diff(self: Series[complex], periods: int = ...) -> Series[complex]: ...  # type: ignore[overload-overlap]
-    @overload
-    def diff(
-        self: Series[bytes] | Series[type] | Series[_str] | Series[Interval],
-        periods: int = ...,
-    ) -> Never: ...
+    def diff(  # type: ignore[overload-overlap]
+        self: Series[Never] | Series[int], periods: int = ...
+    ) -> Series[float]: ...
     @overload
-    def diff(self: Series[Timestamp], periods: int = ...) -> Series[Timedelta]: ...  # type: ignore[overload-overlap]
+    def diff(self: Series[_bool], periods: int = ...) -> Series: ...
     @overload
-    def diff(self: Series[Timedelta], periods: int = ...) -> Series[Timedelta]: ...  # type: ignore[overload-overlap]
+    def diff(self: Series[Period], periods: int = ...) -> OffsetSeries: ...
     @overload
-    def diff(self: Series[Period], periods: int = ...) -> OffsetSeries: ...  # type: ignore[overload-overlap]
+    def diff(self: Series[Interval], periods: int = ...) -> Never: ...
     @overload
-    def diff(self, periods: int = ...) -> Series[float]: ...
+    def diff(
+        self: SupportsGetItem[Scalar, SupportsSelfSub[S1_CO]], periods: int = ...
+    ) -> Series[S1_CO]: ...
     def autocorr(self, lag: int = 1) -> float: ...
     @overload
     def dot(self, other: Series[S1]) -> Scalar: ...
@@ -4665,11 +4669,11 @@ class Series(IndexOpsMixin[S1], NDFrame):
     @overload
     def to_numpy(
         self: Series[Interval],
-        dtype: type[np.bytes_],
+        dtype: type[_T_INTERVAL_NP],
         copy: bool = False,
         na_value: Scalar = ...,
         **kwargs,
-    ) -> np_1darray[np.bytes_]: ...
+    ) -> np_1darray[_T_INTERVAL_NP]: ...
     @overload
     def to_numpy(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,

tests/series/test_series.py

@@ -40,7 +40,6 @@
     Never,
     Self,
     TypeAlias,
-    assert_never,
     assert_type,
 )
 import xarray as xr
@@ -3503,47 +3502,39 @@ def test_diff() -> None:
             BaseOffset,
             index_to_check_for_type=-1,
         )
-    # bool -> object
+    # bool -> Any
     check(
         assert_type(
-            pd.Series([True, True, False, False, True]).diff(),
-            "pd.Series[type[object]]",
+            pd.Series([True, True, False, False, True]).diff(), "pd.Series[Any]"
         ),
         pd.Series,
         object,
     )
-    # object -> object
-    check(
-        assert_type(s.astype(object).diff(), "pd.Series[type[object]]"),
-        pd.Series,
-        object,
-    )
+    # Any -> float
+    s_o = s.astype(object)
+    assert_type(s_o, "pd.Series[Any]")
+    check(assert_type(s_o.diff(), "pd.Series[float]"), pd.Series, float)
     # complex -> complex
     check(
         assert_type(s.astype(complex).diff(), "pd.Series[complex]"), pd.Series, complex
     )
-    if TYPE_CHECKING_INVALID_USAGE:
-        # interval -> TypeError: IntervalArray has no 'diff' method. Convert to a suitable dtype prior to calling 'diff'.
-        assert_never(pd.Series([pd.Interval(0, 2), pd.Interval(1, 4)]).diff())
 
+    def _diff_invalid0():  # pyright: ignore[reportUnusedFunction]
+        # interval -> TypeError: IntervalArray has no 'diff' method. Convert to a suitable dtype prior to calling 'diff'.
+        assert_type(pd.Series([pd.Interval(0, 2), pd.Interval(1, 4)]).diff(), Never)
 
-def test_diff_never1() -> None:
-    s = pd.Series([1, 1, 2, 3, 5, 8])
-    if TYPE_CHECKING_INVALID_USAGE:
+    def _diff_invalid1() -> None:  # pyright: ignore[reportUnusedFunction]
+        s = pd.Series([1, 1, 2, 3, 5, 8])
         # bytes -> numpy.core._exceptions._UFuncNoLoopError: ufunc 'subtract' did not contain a loop with signature matching types (dtype('S21'), dtype('S21')) -> None
-        assert_never(s.astype(bytes).diff())
-
+        s.astype(bytes).diff()  # type: ignore[misc] # pyright: ignore[reportAttributeAccessIssue]
 
-def test_diff_never2() -> None:
-    if TYPE_CHECKING_INVALID_USAGE:
+    def _diff_invalid2() -> None:  # pyright: ignore[reportUnusedFunction]
         # dtype -> TypeError: unsupported operand type(s) for -: 'type' and 'type'
-        assert_never(pd.Series([str, int, bool]).diff())
-
+        pd.Series([str, int, bool]).diff()  # type: ignore[misc] # pyright: ignore[reportAttributeAccessIssue]
 
-def test_diff_never3() -> None:
-    if TYPE_CHECKING_INVALID_USAGE:
+    def _diff_invalid3() -> None:  # pyright: ignore[reportUnusedFunction]
         # str -> TypeError: unsupported operand type(s) for -: 'str' and 'str'
-        assert_never(pd.Series(["a", "b"]).diff())
+        pd.Series(["a", "b"]).diff()  # type: ignore[misc] # pyright: ignore[reportAttributeAccessIssue]
 
 
 def test_operator_constistency() -> None:

tests/test_timefuncs.py

@@ -1063,6 +1063,11 @@ def test_series_types_to_numpy() -> None:
         np_1darray,
         dtype=np.bytes_,
     )
+    check(
+        assert_type(i_s.to_numpy(dtype=np.str_), np_1darray[np.str_]),
+        np_1darray,
+        dtype=np.str_,
+    )
 
 
 def test_index_types_to_numpy() -> None: