make tslibs strptime, timedeltas, and timestamps pass with pyright-strict

pandas-stubs/_libs/tslibs/timedeltas.pyi

@@ -1,3 +1,4 @@
+# pyright: strict
 import datetime as dt
 from datetime import timedelta
 from typing import (
@@ -241,7 +242,9 @@ class Timedelta(timedelta):
     @overload
     def __rmul__(self, other: float) -> Timedelta: ...
     @overload
-    def __rmul__(self, other: np.ndarray) -> np.ndarray: ...
+    def __rmul__(
+        self, other: npt.NDArray[np.floating] | npt.NDArray[np.integer]
+    ) -> npt.NDArray[np.timedelta64]: ...
     @overload
     def __rmul__(self, other: Series[int]) -> TimedeltaSeries: ...
     @overload

pandas-stubs/_libs/tslibs/timestamps.pyi

@@ -1,11 +1,12 @@
+# pyright: strict
 from datetime import (
-    _IsoCalendarDate,
     date as _date,
     datetime,
     time as _time,
     timedelta,
     tzinfo as _tzinfo,
 )
+from datetime import _IsoCalendarDate  # pyright: ignore[reportPrivateUsage]
 import sys
 from time import struct_time
 from typing import (
@@ -15,13 +16,12 @@ from typing import (
     overload,
 )
 
-from _typing import TimeZones
 import numpy as np
 from pandas import (
     DatetimeIndex,
-    Index,
     TimedeltaIndex,
 )
+from pandas.core.indexes.base import UnknownIndex
 from pandas.core.series import (
     Series,
     TimedeltaSeries,
@@ -49,6 +49,9 @@ _Ambiguous: TypeAlias = bool | Literal["raise", "NaT"]
 _Nonexistent: TypeAlias = (
     Literal["raise", "NaT", "shift_backward", "shift_forward"] | Timedelta | timedelta
 )
+# Repeated from `_typing.pyi` so as to satisfy mixed strict / non-strict paths.
+# https://github.com/pandas-dev/pandas-stubs/pull/1151#issuecomment-2715130190
+TimeZones: TypeAlias = str | _tzinfo | None | int
 
 class Timestamp(datetime, SupportsIndex):
     min: ClassVar[Timestamp]  # pyright: ignore[reportIncompatibleVariableOverride]
@@ -166,25 +169,33 @@ class Timestamp(datetime, SupportsIndex):
     @overload  # type: ignore[override]
     def __le__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
     @overload
-    def __le__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ...
+    def __le__(
+        self, other: DatetimeIndex | npt.NDArray[np.datetime64]
+    ) -> np_ndarray_bool: ...
     @overload
     def __le__(self, other: TimestampSeries) -> Series[bool]: ...
     @overload  # type: ignore[override]
     def __lt__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
     @overload
-    def __lt__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ...
+    def __lt__(
+        self, other: DatetimeIndex | npt.NDArray[np.datetime64]
+    ) -> np_ndarray_bool: ...
     @overload
     def __lt__(self, other: TimestampSeries) -> Series[bool]: ...
     @overload  # type: ignore[override]
     def __ge__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
     @overload
-    def __ge__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ...
+    def __ge__(
+        self, other: DatetimeIndex | npt.NDArray[np.datetime64]
+    ) -> np_ndarray_bool: ...
     @overload
     def __ge__(self, other: TimestampSeries) -> Series[bool]: ...
     @overload  # type: ignore[override]
     def __gt__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
     @overload
-    def __gt__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ...
+    def __gt__(
+        self, other: DatetimeIndex | npt.NDArray[np.datetime64]
+    ) -> np_ndarray_bool: ...
     @overload
     def __gt__(self, other: TimestampSeries) -> Series[bool]: ...
     # error: Signature of "__add__" incompatible with supertype "date"/"datetime"
@@ -224,15 +235,15 @@ class Timestamp(datetime, SupportsIndex):
     @overload
     def __eq__(self, other: TimestampSeries) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload
-    def __eq__(self, other: npt.NDArray[np.datetime64] | Index) -> np_ndarray_bool: ...  # type: ignore[overload-overlap]
+    def __eq__(self, other: npt.NDArray[np.datetime64] | UnknownIndex) -> np_ndarray_bool: ...  # type: ignore[overload-overlap]
     @overload
     def __eq__(self, other: object) -> Literal[False]: ...
     @overload
     def __ne__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     @overload
     def __ne__(self, other: TimestampSeries) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload
-    def __ne__(self, other: npt.NDArray[np.datetime64] | Index) -> np_ndarray_bool: ...  # type: ignore[overload-overlap]
+    def __ne__(self, other: npt.NDArray[np.datetime64] | UnknownIndex) -> np_ndarray_bool: ...  # type: ignore[overload-overlap]
     @overload
     def __ne__(self, other: object) -> Literal[True]: ...
     def __hash__(self) -> int: ...

pandas-stubs/_typing.pyi

@@ -833,6 +833,8 @@ TimeGrouperOrigin: TypeAlias = (
 ExcelReadEngine: TypeAlias = Literal["xlrd", "openpyxl", "odf", "pyxlsb", "calamine"]
 ExcelWriteEngine: TypeAlias = Literal["openpyxl", "odf", "xlsxwriter"]
 
+# Repeated in `timestamps.pyi` so as to satisfy mixed strict / non-strict paths.
+# https://github.com/pandas-dev/pandas-stubs/pull/1151#issuecomment-2715130190
 TimeZones: TypeAlias = str | tzinfo | None | int
 
 __all__ = ["npt", "type_t"]

pandas-stubs/core/frame.pyi

@@ -31,7 +31,10 @@ from pandas.core.generic import NDFrame
 from pandas.core.groupby.generic import DataFrameGroupBy
 from pandas.core.groupby.grouper import Grouper
 from pandas.core.indexers import BaseIndexer
-from pandas.core.indexes.base import Index
+from pandas.core.indexes.base import (
+    Index,
+    UnknownIndex,
+)
 from pandas.core.indexes.category import CategoricalIndex
 from pandas.core.indexes.datetimes import DatetimeIndex
 from pandas.core.indexes.interval import IntervalIndex
@@ -1617,7 +1620,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     @property
     def iloc(self) -> _iLocIndexerFrame[Self]: ...
     @property
-    def index(self) -> Index: ...
+    def index(self) -> UnknownIndex: ...
     @index.setter
     def index(self, idx: Index) -> None: ...
     @property

tests/test_scalars.py

@@ -713,8 +713,16 @@ def test_timedelta_mul_div() -> None:
 
     check(assert_type(md_int * td, pd.Timedelta), pd.Timedelta)
     check(assert_type(md_float * td, pd.Timedelta), pd.Timedelta)
-    check(assert_type(md_ndarray_intp * td, np.ndarray), np.ndarray, np.timedelta64)
-    check(assert_type(md_ndarray_float * td, np.ndarray), np.ndarray, np.timedelta64)
+    check(
+        assert_type(md_ndarray_intp * td, npt.NDArray[np.timedelta64]),
+        np.ndarray,
+        np.timedelta64,
+    )
+    check(
+        assert_type(md_ndarray_float * td, npt.NDArray[np.timedelta64]),
+        np.ndarray,
+        np.timedelta64,
+    )
     check(assert_type(mp_series_int * td, TimedeltaSeries), pd.Series, pd.Timedelta)
     check(assert_type(md_series_float * td, TimedeltaSeries), pd.Series, pd.Timedelta)
     check(assert_type(md_int64_index * td, pd.TimedeltaIndex), pd.TimedeltaIndex)
@@ -1253,6 +1261,8 @@ def test_timestamp_cmp() -> None:
     c_np_dt64 = np.datetime64(1, "ns")
     c_dt_datetime = dt.datetime(year=2000, month=1, day=1)
     c_datetimeindex = pd.DatetimeIndex(["2000-1-1"])
+    # DatetimeIndex, but the type checker detects it to be UnknownIndex.
+    c_unknown_index = pd.DataFrame({"a": [1]}, index=c_datetimeindex).index
     c_np_ndarray_dt64 = np_dt64_arr
     c_series_dt64: TimestampSeries = pd.Series([1, 2, 3], dtype="datetime64[ns]")
     c_series_timestamp = pd.Series(pd.DatetimeIndex(["2000-1-1"]))
@@ -1273,6 +1283,8 @@ def test_timestamp_cmp() -> None:
 
     check(assert_type(ts > c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_)
     check(assert_type(ts <= c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_)
+    check(assert_type(ts > c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_)
+    check(assert_type(ts <= c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_)
 
     check(assert_type(ts > c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_)
     check(assert_type(ts <= c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_)
@@ -1292,6 +1304,8 @@ def test_timestamp_cmp() -> None:
 
     check(assert_type(c_datetimeindex > ts, np_ndarray_bool), np.ndarray, np.bool_)
     check(assert_type(c_datetimeindex <= ts, np_ndarray_bool), np.ndarray, np.bool_)
+    check(assert_type(c_unknown_index > ts, np_ndarray_bool), np.ndarray, np.bool_)
+    check(assert_type(c_unknown_index <= ts, np_ndarray_bool), np.ndarray, np.bool_)
 
     check(assert_type(c_np_ndarray_dt64 > ts, np_ndarray_bool), np.ndarray, np.bool_)
     check(assert_type(c_np_ndarray_dt64 <= ts, np_ndarray_bool), np.ndarray, np.bool_)
@@ -1313,6 +1327,8 @@ def test_timestamp_cmp() -> None:
 
     check(assert_type(ts >= c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_)
     check(assert_type(ts < c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_)
+    check(assert_type(ts >= c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_)
+    check(assert_type(ts < c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_)
 
     check(assert_type(ts >= c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_)
     check(assert_type(ts < c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_)
@@ -1332,6 +1348,8 @@ def test_timestamp_cmp() -> None:
 
     check(assert_type(c_datetimeindex >= ts, np_ndarray_bool), np.ndarray, np.bool_)
     check(assert_type(c_datetimeindex < ts, np_ndarray_bool), np.ndarray, np.bool_)
+    check(assert_type(c_unknown_index >= ts, np_ndarray_bool), np.ndarray, np.bool_)
+    check(assert_type(c_unknown_index < ts, np_ndarray_bool), np.ndarray, np.bool_)
 
     check(assert_type(c_np_ndarray_dt64 >= ts, np_ndarray_bool), np.ndarray, np.bool_)
     check(assert_type(c_np_ndarray_dt64 < ts, np_ndarray_bool), np.ndarray, np.bool_)
@@ -1358,6 +1376,13 @@ def test_timestamp_cmp() -> None:
         assert_type(ts != c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_
     )
     assert (eq_arr != ne_arr).all()
+    eq_arr = check(
+        assert_type(ts == c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_
+    )
+    ne_arr = check(
+        assert_type(ts != c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_
+    )
+    assert (eq_arr != ne_arr).all()
 
     eq_arr = check(
         assert_type(ts == c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_
@@ -1396,6 +1421,8 @@ def test_timestamp_eq_ne_rhs() -> None:
     c_np_dt64 = np.datetime64(1, "ns")
     c_dt_datetime = dt.datetime(year=2000, month=1, day=1)
     c_datetimeindex = pd.DatetimeIndex(["2000-1-1"])
+    # DatetimeIndex, but the type checker detects it to be UnknownIndex.
+    c_unknown_index = pd.DataFrame({"a": [1]}, index=c_datetimeindex).index
     c_np_ndarray_dt64 = np_dt64_arr
     c_series_dt64: pd.Series[pd.Timestamp] = pd.Series(
         [1, 2, 3], dtype="datetime64[ns]"
@@ -1416,6 +1443,13 @@ def test_timestamp_eq_ne_rhs() -> None:
         assert_type(c_datetimeindex != ts, np_ndarray_bool), np.ndarray, np.bool_
     )
     assert (eq_arr != ne_arr).all()
+    eq_arr = check(
+        assert_type(c_unknown_index == ts, np_ndarray_bool), np.ndarray, np.bool_
+    )
+    ne_arr = check(
+        assert_type(c_unknown_index != ts, np_ndarray_bool), np.ndarray, np.bool_
+    )
+    assert (eq_arr != ne_arr).all()
 
     eq_a = check(assert_type(c_np_ndarray_dt64 != ts, Any), np.ndarray, np.bool_)
     ne_a = check(assert_type(c_np_ndarray_dt64 == ts, Any), np.ndarray, np.bool_)