Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
Mar 12, 2025
Merged

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

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
a2a8cd1
make some tslibs files pass with pyright-strict
MarcoGorelli Mar 7, 2025
4bf8eb7
Merge remote-tracking branch 'upstream/main' into strict-typing-tslibs-i
MarcoGorelli Mar 10, 2025
5fb4fcb
add test which uses unknown index
MarcoGorelli Mar 10, 2025
f39c5c4
try reverting strict in timestamps.pyi
MarcoGorelli Mar 10, 2025
ae12843
use Index[Any] for __eq__ and __neq__
MarcoGorelli Mar 10, 2025
e39056f
try putting strict back
MarcoGorelli Mar 10, 2025
c2d69d3
Merge remote-tracking branch 'upstream/main' into strict-typing-tslibs-i
MarcoGorelli Mar 11, 2025
21c6f6b
DataFrame.index -> UnknownIndex
MarcoGorelli Mar 11, 2025
707c988
copy TimeZones over
MarcoGorelli Mar 11, 2025
2a7349a
comment link
MarcoGorelli Mar 11, 2025
b05046e
weve got unknownindex, lets use it
MarcoGorelli Mar 11, 2025
cb99f48
extra tests, clean up
MarcoGorelli Mar 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pandas-stubs/_libs/tslibs/strptime.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
# pyright
from typing import Any

def array_strptime(*args, **kwargs) -> Any: ...
5 changes: 4 additions & 1 deletion pandas-stubs/_libs/tslibs/timedeltas.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
# pyright: strict
import datetime as dt
from datetime import timedelta
from typing import (
Expand Down Expand Up @@ -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
Expand Down
29 changes: 20 additions & 9 deletions pandas-stubs/_libs/tslibs/timestamps.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 (
Expand All @@ -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,
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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: ...
Expand Down
7 changes: 5 additions & 2 deletions pandas-stubs/core/frame.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
25 changes: 23 additions & 2 deletions tests/test_scalars.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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 Index[Unknown].
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"]))
Expand All @@ -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_)
Expand Down Expand Up @@ -1313,6 +1325,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_)
Expand Down Expand Up @@ -1358,6 +1372,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_
Expand Down