diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 081f11e9e..01a17e1df 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/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 diff --git a/pandas-stubs/_libs/tslibs/timestamps.pyi b/pandas-stubs/_libs/tslibs/timestamps.pyi index 702c1287d..8cae60fec 100644 --- a/pandas-stubs/_libs/tslibs/timestamps.pyi +++ b/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,7 +235,7 @@ 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 @@ -232,7 +243,7 @@ class Timestamp(datetime, SupportsIndex): @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: ... diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 478f60da0..46c2c6438 100644 --- a/pandas-stubs/_typing.pyi +++ b/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"] diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 03fab417e..0fd4a673e 100644 --- a/pandas-stubs/core/frame.pyi +++ b/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 diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 69a3dbf7d..0b130320b 100644 --- a/tests/test_scalars.py +++ b/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_)