Skip to content

Commit af65974

Browse files
loicdiridollouloicdiridollou
committed
Merge branch 'main' of github.com:loicdiridollou/pandas-stubs into gh1089_fix_nightly
2 parents 5050ffe + f6f7bb0 commit af65974

File tree

7 files changed

+98
-23
lines changed

7 files changed

+98
-23
lines changed

‎pandas-stubs/_libs/tslibs/timedeltas.pyi‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
# pyright: strict
12
import datetime as dt
23
from datetime import timedelta
34
from typing import (
@@ -241,7 +242,9 @@ class Timedelta(timedelta):
241242
@overload
242243
def __rmul__(self, other: float) -> Timedelta: ...
243244
@overload
244-
def __rmul__(self, other: np.ndarray) -> np.ndarray: ...
245+
def __rmul__(
246+
self, other: npt.NDArray[np.floating] | npt.NDArray[np.integer]
247+
) -> npt.NDArray[np.timedelta64]: ...
245248
@overload
246249
def __rmul__(self, other: Series[int]) -> TimedeltaSeries: ...
247250
@overload

‎pandas-stubs/_libs/tslibs/timestamps.pyi‎

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
1+
# pyright: strict
12
from datetime import (
2-
_IsoCalendarDate,
33
date as _date,
44
datetime,
55
time as _time,
66
timedelta,
77
tzinfo as _tzinfo,
88
)
9+
from datetime import _IsoCalendarDate # pyright: ignore[reportPrivateUsage]
910
import sys
1011
from time import struct_time
1112
from typing import (
@@ -15,13 +16,12 @@ from typing import (
1516
overload,
1617
)
1718

18-
from _typing import TimeZones
1919
import numpy as np
2020
from pandas import (
2121
DatetimeIndex,
22-
Index,
2322
TimedeltaIndex,
2423
)
24+
from pandas.core.indexes.base import UnknownIndex
2525
from pandas.core.series import (
2626
Series,
2727
TimedeltaSeries,
@@ -49,6 +49,9 @@ _Ambiguous: TypeAlias = bool | Literal["raise", "NaT"]
4949
_Nonexistent: TypeAlias = (
5050
Literal["raise", "NaT", "shift_backward", "shift_forward"] | Timedelta | timedelta
5151
)
52+
# Repeated from `_typing.pyi` so as to satisfy mixed strict / non-strict paths.
53+
# https://github.com/pandas-dev/pandas-stubs/pull/1151#issuecomment-2715130190
54+
TimeZones: TypeAlias = str | _tzinfo | None | int
5255

5356
class Timestamp(datetime, SupportsIndex):
5457
min: ClassVar[Timestamp] # pyright: ignore[reportIncompatibleVariableOverride]
@@ -166,25 +169,33 @@ class Timestamp(datetime, SupportsIndex):
166169
@overload # type: ignore[override]
167170
def __le__(self, other: Timestamp | datetime | np.datetime64) -> bool: ... # type: ignore[misc]
168171
@overload
169-
def __le__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ...
172+
def __le__(
173+
self, other: DatetimeIndex | npt.NDArray[np.datetime64]
174+
) -> np_ndarray_bool: ...
170175
@overload
171176
def __le__(self, other: TimestampSeries) -> Series[bool]: ...
172177
@overload # type: ignore[override]
173178
def __lt__(self, other: Timestamp | datetime | np.datetime64) -> bool: ... # type: ignore[misc]
174179
@overload
175-
def __lt__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ...
180+
def __lt__(
181+
self, other: DatetimeIndex | npt.NDArray[np.datetime64]
182+
) -> np_ndarray_bool: ...
176183
@overload
177184
def __lt__(self, other: TimestampSeries) -> Series[bool]: ...
178185
@overload # type: ignore[override]
179186
def __ge__(self, other: Timestamp | datetime | np.datetime64) -> bool: ... # type: ignore[misc]
180187
@overload
181-
def __ge__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ...
188+
def __ge__(
189+
self, other: DatetimeIndex | npt.NDArray[np.datetime64]
190+
) -> np_ndarray_bool: ...
182191
@overload
183192
def __ge__(self, other: TimestampSeries) -> Series[bool]: ...
184193
@overload # type: ignore[override]
185194
def __gt__(self, other: Timestamp | datetime | np.datetime64) -> bool: ... # type: ignore[misc]
186195
@overload
187-
def __gt__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ...
196+
def __gt__(
197+
self, other: DatetimeIndex | npt.NDArray[np.datetime64]
198+
) -> np_ndarray_bool: ...
188199
@overload
189200
def __gt__(self, other: TimestampSeries) -> Series[bool]: ...
190201
# error: Signature of "__add__" incompatible with supertype "date"/"datetime"
@@ -224,15 +235,15 @@ class Timestamp(datetime, SupportsIndex):
224235
@overload
225236
def __eq__(self, other: TimestampSeries) -> Series[bool]: ... # type: ignore[overload-overlap]
226237
@overload
227-
def __eq__(self, other: npt.NDArray[np.datetime64] | Index) -> np_ndarray_bool: ... # type: ignore[overload-overlap]
238+
def __eq__(self, other: npt.NDArray[np.datetime64] | UnknownIndex) -> np_ndarray_bool: ... # type: ignore[overload-overlap]
228239
@overload
229240
def __eq__(self, other: object) -> Literal[False]: ...
230241
@overload
231242
def __ne__(self, other: Timestamp | datetime | np.datetime64) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
232243
@overload
233244
def __ne__(self, other: TimestampSeries) -> Series[bool]: ... # type: ignore[overload-overlap]
234245
@overload
235-
def __ne__(self, other: npt.NDArray[np.datetime64] | Index) -> np_ndarray_bool: ... # type: ignore[overload-overlap]
246+
def __ne__(self, other: npt.NDArray[np.datetime64] | UnknownIndex) -> np_ndarray_bool: ... # type: ignore[overload-overlap]
236247
@overload
237248
def __ne__(self, other: object) -> Literal[True]: ...
238249
def __hash__(self) -> int: ...

‎pandas-stubs/_typing.pyi‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -740,7 +740,7 @@ ReplaceValue: TypeAlias = (
740740
| Pattern
741741
| NAType
742742
| Sequence[Scalar | Pattern]
743-
| Mapping[Hashable, Scalar]
743+
| Mapping[HashableT, ScalarT]
744744
| Series[Any]
745745
| None
746746
)
@@ -833,6 +833,8 @@ TimeGrouperOrigin: TypeAlias = (
833833
ExcelReadEngine: TypeAlias = Literal["xlrd", "openpyxl", "odf", "pyxlsb", "calamine"]
834834
ExcelWriteEngine: TypeAlias = Literal["openpyxl", "odf", "xlsxwriter"]
835835

836+
# Repeated in `timestamps.pyi` so as to satisfy mixed strict / non-strict paths.
837+
# https://github.com/pandas-dev/pandas-stubs/pull/1151#issuecomment-2715130190
836838
TimeZones: TypeAlias = str | tzinfo | None | int
837839

838840
__all__ = ["npt", "type_t"]

‎pandas-stubs/core/frame.pyi‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,10 @@ from pandas.core.generic import NDFrame
3131
from pandas.core.groupby.generic import DataFrameGroupBy
3232
from pandas.core.groupby.grouper import Grouper
3333
from pandas.core.indexers import BaseIndexer
34-
from pandas.core.indexes.base import Index
34+
from pandas.core.indexes.base import (
35+
Index,
36+
UnknownIndex,
37+
)
3538
from pandas.core.indexes.category import CategoricalIndex
3639
from pandas.core.indexes.datetimes import DatetimeIndex
3740
from pandas.core.indexes.interval import IntervalIndex
@@ -89,6 +92,7 @@ from pandas._typing import (
8992
HashableT,
9093
HashableT1,
9194
HashableT2,
95+
HashableT3,
9296
IgnoreRaise,
9397
IndexingInt,
9498
IndexLabel,
@@ -181,7 +185,7 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]):
181185
IndexType
182186
| MaskType
183187
| Callable[[DataFrame], IndexType | MaskType | Sequence[Hashable]]
184-
| list[Hashable]
188+
| list[HashableT]
185189
| tuple[
186190
IndexType
187191
| MaskType
@@ -794,20 +798,20 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
794798
@overload
795799
def replace(
796800
self,
797-
to_replace: ReplaceValue | Mapping[Hashable, ReplaceValue] = ...,
798-
value: ReplaceValue | Mapping[Hashable, ReplaceValue] = ...,
801+
to_replace: ReplaceValue | Mapping[HashableT2, ReplaceValue] = ...,
802+
value: ReplaceValue | Mapping[HashableT3, ReplaceValue] = ...,
799803
*,
800804
inplace: Literal[True],
801-
regex: ReplaceValue | Mapping[Hashable, ReplaceValue] = ...,
805+
regex: ReplaceValue | Mapping[HashableT3, ReplaceValue] = ...,
802806
) -> None: ...
803807
@overload
804808
def replace(
805809
self,
806-
to_replace: ReplaceValue | Mapping[Hashable, ReplaceValue] = ...,
807-
value: ReplaceValue | Mapping[Hashable, ReplaceValue] = ...,
810+
to_replace: ReplaceValue | Mapping[HashableT2, ReplaceValue] = ...,
811+
value: ReplaceValue | Mapping[HashableT3, ReplaceValue] = ...,
808812
*,
809813
inplace: Literal[False] = ...,
810-
regex: ReplaceValue | Mapping[Hashable, ReplaceValue] = ...,
814+
regex: ReplaceValue | Mapping[HashableT3, ReplaceValue] = ...,
811815
) -> Self: ...
812816
def shift(
813817
self,
@@ -1617,7 +1621,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
16171621
@property
16181622
def iloc(self) -> _iLocIndexerFrame[Self]: ...
16191623
@property
1620-
def index(self) -> Index: ...
1624+
def index(self) -> UnknownIndex: ...
16211625
@index.setter
16221626
def index(self, idx: Index) -> None: ...
16231627
@property

‎tests/test_frame.py‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2574,8 +2574,10 @@ def test_types_replace() -> None:
25742574

25752575

25762576
def test_dataframe_replace() -> None:
2577-
df = pd.DataFrame({"col1": ["a", "ab", "ba"]})
2577+
df = pd.DataFrame({"col1": ["a", "ab", "ba"], "col2": [0, 1, 2]})
25782578
pattern = re.compile(r"^a.*")
2579+
replace_dict_scalar = {0: 1}
2580+
replace_dict_per_column = {"col2": {0: 1}}
25792581
check(assert_type(df.replace("a", "x"), pd.DataFrame), pd.DataFrame)
25802582
check(assert_type(df.replace(pattern, "x"), pd.DataFrame), pd.DataFrame)
25812583
check(assert_type(df.replace("a", "x", regex=True), pd.DataFrame), pd.DataFrame)
@@ -2592,6 +2594,7 @@ def test_dataframe_replace() -> None:
25922594
)
25932595

25942596
check(assert_type(df.replace({"a": "x"}), pd.DataFrame), pd.DataFrame)
2597+
check(assert_type(df.replace(replace_dict_scalar), pd.DataFrame), pd.DataFrame)
25952598
check(assert_type(df.replace({pattern: "x"}), pd.DataFrame), pd.DataFrame)
25962599
check(assert_type(df.replace(pd.Series({"a": "x"})), pd.DataFrame), pd.DataFrame)
25972600
check(assert_type(df.replace(regex={"a": "x"}), pd.DataFrame), pd.DataFrame)
@@ -2670,6 +2673,7 @@ def test_dataframe_replace() -> None:
26702673
)
26712674

26722675
check(assert_type(df.replace({"col1": {"a": "x"}}), pd.DataFrame), pd.DataFrame)
2676+
check(assert_type(df.replace(replace_dict_per_column), pd.DataFrame), pd.DataFrame)
26732677
check(assert_type(df.replace({"col1": {pattern: "x"}}), pd.DataFrame), pd.DataFrame)
26742678
check(
26752679
assert_type(df.replace({"col1": pd.Series({"a": "x"})}), pd.DataFrame),
@@ -3582,6 +3586,18 @@ def test_in_columns() -> None:
35823586
check(assert_type(df.groupby(by=cols).sum(), pd.DataFrame), pd.DataFrame)
35833587

35843588

3589+
def test_loc_list_str() -> None:
3590+
# GH 1162 (PR)
3591+
df = pd.DataFrame(
3592+
[[1, 2], [4, 5], [7, 8]],
3593+
index=["cobra", "viper", "sidewinder"],
3594+
columns=["max_speed", "shield"],
3595+
)
3596+
3597+
result = df.loc[["viper", "sidewinder"]]
3598+
check(assert_type(result, pd.DataFrame), pd.DataFrame)
3599+
3600+
35853601
def test_insert_newvalues() -> None:
35863602
df = pd.DataFrame({"a": [1, 2]})
35873603
ab = pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})

‎tests/test_scalars.py‎

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -713,8 +713,16 @@ def test_timedelta_mul_div() -> None:
713713

714714
check(assert_type(md_int * td, pd.Timedelta), pd.Timedelta)
715715
check(assert_type(md_float * td, pd.Timedelta), pd.Timedelta)
716-
check(assert_type(md_ndarray_intp * td, np.ndarray), np.ndarray, np.timedelta64)
717-
check(assert_type(md_ndarray_float * td, np.ndarray), np.ndarray, np.timedelta64)
716+
check(
717+
assert_type(md_ndarray_intp * td, npt.NDArray[np.timedelta64]),
718+
np.ndarray,
719+
np.timedelta64,
720+
)
721+
check(
722+
assert_type(md_ndarray_float * td, npt.NDArray[np.timedelta64]),
723+
np.ndarray,
724+
np.timedelta64,
725+
)
718726
check(assert_type(mp_series_int * td, TimedeltaSeries), pd.Series, pd.Timedelta)
719727
check(assert_type(md_series_float * td, TimedeltaSeries), pd.Series, pd.Timedelta)
720728
check(assert_type(md_int64_index * td, pd.TimedeltaIndex), pd.TimedeltaIndex)
@@ -1253,6 +1261,8 @@ def test_timestamp_cmp() -> None:
12531261
c_np_dt64 = np.datetime64(1, "ns")
12541262
c_dt_datetime = dt.datetime(year=2000, month=1, day=1)
12551263
c_datetimeindex = pd.DatetimeIndex(["2000-1-1"])
1264+
# DatetimeIndex, but the type checker detects it to be UnknownIndex.
1265+
c_unknown_index = pd.DataFrame({"a": [1]}, index=c_datetimeindex).index
12561266
c_np_ndarray_dt64 = np_dt64_arr
12571267
c_series_dt64: TimestampSeries = pd.Series([1, 2, 3], dtype="datetime64[ns]")
12581268
c_series_timestamp = pd.Series(pd.DatetimeIndex(["2000-1-1"]))
@@ -1273,6 +1283,8 @@ def test_timestamp_cmp() -> None:
12731283

12741284
check(assert_type(ts > c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_)
12751285
check(assert_type(ts <= c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_)
1286+
check(assert_type(ts > c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_)
1287+
check(assert_type(ts <= c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_)
12761288

12771289
check(assert_type(ts > c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_)
12781290
check(assert_type(ts <= c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_)
@@ -1292,6 +1304,8 @@ def test_timestamp_cmp() -> None:
12921304

12931305
check(assert_type(c_datetimeindex > ts, np_ndarray_bool), np.ndarray, np.bool_)
12941306
check(assert_type(c_datetimeindex <= ts, np_ndarray_bool), np.ndarray, np.bool_)
1307+
check(assert_type(c_unknown_index > ts, np_ndarray_bool), np.ndarray, np.bool_)
1308+
check(assert_type(c_unknown_index <= ts, np_ndarray_bool), np.ndarray, np.bool_)
12951309

12961310
check(assert_type(c_np_ndarray_dt64 > ts, np_ndarray_bool), np.ndarray, np.bool_)
12971311
check(assert_type(c_np_ndarray_dt64 <= ts, np_ndarray_bool), np.ndarray, np.bool_)
@@ -1313,6 +1327,8 @@ def test_timestamp_cmp() -> None:
13131327

13141328
check(assert_type(ts >= c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_)
13151329
check(assert_type(ts < c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_)
1330+
check(assert_type(ts >= c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_)
1331+
check(assert_type(ts < c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_)
13161332

13171333
check(assert_type(ts >= c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_)
13181334
check(assert_type(ts < c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_)
@@ -1332,6 +1348,8 @@ def test_timestamp_cmp() -> None:
13321348

13331349
check(assert_type(c_datetimeindex >= ts, np_ndarray_bool), np.ndarray, np.bool_)
13341350
check(assert_type(c_datetimeindex < ts, np_ndarray_bool), np.ndarray, np.bool_)
1351+
check(assert_type(c_unknown_index >= ts, np_ndarray_bool), np.ndarray, np.bool_)
1352+
check(assert_type(c_unknown_index < ts, np_ndarray_bool), np.ndarray, np.bool_)
13351353

13361354
check(assert_type(c_np_ndarray_dt64 >= ts, np_ndarray_bool), np.ndarray, np.bool_)
13371355
check(assert_type(c_np_ndarray_dt64 < ts, np_ndarray_bool), np.ndarray, np.bool_)
@@ -1358,6 +1376,13 @@ def test_timestamp_cmp() -> None:
13581376
assert_type(ts != c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_
13591377
)
13601378
assert (eq_arr != ne_arr).all()
1379+
eq_arr = check(
1380+
assert_type(ts == c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_
1381+
)
1382+
ne_arr = check(
1383+
assert_type(ts != c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_
1384+
)
1385+
assert (eq_arr != ne_arr).all()
13611386

13621387
eq_arr = check(
13631388
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:
13961421
c_np_dt64 = np.datetime64(1, "ns")
13971422
c_dt_datetime = dt.datetime(year=2000, month=1, day=1)
13981423
c_datetimeindex = pd.DatetimeIndex(["2000-1-1"])
1424+
# DatetimeIndex, but the type checker detects it to be UnknownIndex.
1425+
c_unknown_index = pd.DataFrame({"a": [1]}, index=c_datetimeindex).index
13991426
c_np_ndarray_dt64 = np_dt64_arr
14001427
c_series_dt64: pd.Series[pd.Timestamp] = pd.Series(
14011428
[1, 2, 3], dtype="datetime64[ns]"
@@ -1416,6 +1443,13 @@ def test_timestamp_eq_ne_rhs() -> None:
14161443
assert_type(c_datetimeindex != ts, np_ndarray_bool), np.ndarray, np.bool_
14171444
)
14181445
assert (eq_arr != ne_arr).all()
1446+
eq_arr = check(
1447+
assert_type(c_unknown_index == ts, np_ndarray_bool), np.ndarray, np.bool_
1448+
)
1449+
ne_arr = check(
1450+
assert_type(c_unknown_index != ts, np_ndarray_bool), np.ndarray, np.bool_
1451+
)
1452+
assert (eq_arr != ne_arr).all()
14191453

14201454
eq_a = check(assert_type(c_np_ndarray_dt64 != ts, Any), np.ndarray, np.bool_)
14211455
ne_a = check(assert_type(c_np_ndarray_dt64 == ts, Any), np.ndarray, np.bool_)

‎tests/test_series.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1448,12 +1448,17 @@ def test_types_replace() -> None:
14481448
def test_series_replace() -> None:
14491449
s: pd.Series[str] = pd.DataFrame({"col1": ["a", "ab", "ba"]})["col1"]
14501450
pattern = re.compile(r"^a.*")
1451+
replace_dict = {"a": "b"}
14511452
check(assert_type(s.replace("a", "x"), "pd.Series[str]"), pd.Series)
14521453
check(assert_type(s.replace(pattern, "x"), "pd.Series[str]"), pd.Series)
14531454
check(
14541455
assert_type(s.replace({"a": "z"}), "pd.Series[str]"),
14551456
pd.Series,
14561457
)
1458+
check(
1459+
assert_type(s.replace(replace_dict), "pd.Series[str]"),
1460+
pd.Series,
1461+
)
14571462
check(
14581463
assert_type(s.replace(pd.Series({"a": "z"})), "pd.Series[str]"),
14591464
pd.Series,

0 commit comments

Comments
 (0)