Skip to content

Commit f7621f4

Browse files
twoertweinDr-Irv
twoertwein
and
Dr-Irv
authored
Make IndexOpsMixin (and Index) generic (#760)
* interpolate method * WIP * mypy handles Never differently - at least asssert that the method never returns * use Self * a few more S1/Self * fix tests * more tests * add an ignore to let mypy pass * compat with old python versions * fix CategoricalIndex; MultiIndex should probably be Index[tuple[S1, ...]] but keeping it as Index[Any] * Revert "interpolate method" This reverts commit 4929ecb. * many more overloads for subclasses (I wish pandas would not handle subclasses in parent classes) * remove PandasObject - it caused the pyright issues but it provides no methods that aren't already provided by object * works (except for np.ndarray) * address the easy-to-fix comments * overloads for numpy * did numexpr break the CI? * lie: df.columns -> Index[str] * new ruff has far more pyi rules * Make it clear that both mypy&pyright infer Never; but only pyright detects it as invalid * use type aliases instead of npt * fix issue with floordiv for mypy * fix comment in arraylike --------- Co-authored-by: Irv Lustig <[email protected]>
1 parent fbec52b commit f7621f4

35 files changed

+540
-431
lines changed

‎.pre-commit-config.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ repos:
1111
hooks:
1212
- id: isort
1313
- repo: https://github.com/astral-sh/ruff-pre-commit
14-
rev: v0.0.282
14+
rev: v0.0.283
1515
hooks:
1616
- id: ruff
1717
args: [

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

+7-11
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,11 @@ import numpy as np
1111
import pandas as pd
1212
from pandas import (
1313
DatetimeIndex,
14+
Index,
1415
PeriodIndex,
1516
Series,
1617
TimedeltaIndex,
1718
)
18-
from pandas.core.indexes.base import (
19-
_FloatIndexType,
20-
_IntIndexType,
21-
)
2219
from pandas.core.series import (
2320
TimedeltaSeries,
2421
TimestampSeries,
@@ -245,7 +242,7 @@ class Timedelta(timedelta):
245242
@overload
246243
def __mul__(self, other: Series[float]) -> TimedeltaSeries: ...
247244
@overload
248-
def __mul__(self, other: _IntIndexType | _FloatIndexType) -> TimedeltaIndex: ...
245+
def __mul__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
249246
@overload
250247
def __rmul__(self, other: float) -> Timedelta: ...
251248
@overload
@@ -254,8 +251,9 @@ class Timedelta(timedelta):
254251
def __rmul__(self, other: Series[int]) -> TimedeltaSeries: ...
255252
@overload
256253
def __rmul__(self, other: Series[float]) -> TimedeltaSeries: ...
254+
# maybe related to https://github.com/python/mypy/issues/10755
257255
@overload
258-
def __rmul__(self, other: _IntIndexType | _FloatIndexType) -> TimedeltaIndex: ...
256+
def __rmul__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ... # type: ignore[misc]
259257
# Override due to more types supported than dt.timedelta
260258
# error: Signature of "__floordiv__" incompatible with supertype "timedelta"
261259
@overload # type: ignore[override]
@@ -271,9 +269,7 @@ class Timedelta(timedelta):
271269
self, other: npt.NDArray[np.timedelta64]
272270
) -> npt.NDArray[np.int_]: ...
273271
@overload
274-
def __floordiv__(
275-
self, other: _IntIndexType | _FloatIndexType
276-
) -> TimedeltaIndex: ...
272+
def __floordiv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
277273
@overload
278274
def __floordiv__(self, other: Series[int]) -> TimedeltaSeries: ...
279275
@overload
@@ -306,7 +302,7 @@ class Timedelta(timedelta):
306302
@overload
307303
def __truediv__(self, other: Series[float]) -> TimedeltaSeries: ...
308304
@overload
309-
def __truediv__(self, other: _IntIndexType | _FloatIndexType) -> TimedeltaIndex: ...
305+
def __truediv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
310306
def __rtruediv__(self, other: timedelta | Timedelta | NaTType) -> float: ...
311307
# Override due to more types supported than dt.timedelta
312308
@overload
@@ -338,7 +334,7 @@ class Timedelta(timedelta):
338334
@overload
339335
def __mod__(self, other: Series[int] | Series[float]) -> TimedeltaSeries: ...
340336
@overload
341-
def __mod__(self, other: _IntIndexType | _FloatIndexType) -> TimedeltaIndex: ...
337+
def __mod__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
342338
@overload
343339
def __mod__(
344340
self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]

‎pandas-stubs/_typing.pyi

+17-30
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ DtypeBackend: TypeAlias = Literal["pyarrow", "numpy_nullable"]
9191

9292
BooleanDtypeArg: TypeAlias = (
9393
# Builtin bool type and its string alias
94-
type[bool] # noqa: PYI030
94+
type[bool] # noqa: PYI030,PYI055
9595
| Literal["bool"]
9696
# Pandas nullable boolean type and its string alias
9797
| pd.BooleanDtype
@@ -105,7 +105,7 @@ BooleanDtypeArg: TypeAlias = (
105105
)
106106
IntDtypeArg: TypeAlias = (
107107
# Builtin integer type and its string alias
108-
type[int] # noqa: PYI030
108+
type[int] # noqa: PYI030,PYI055
109109
| Literal["int"]
110110
# Pandas nullable integer types and their string aliases
111111
| pd.Int8Dtype
@@ -137,7 +137,7 @@ IntDtypeArg: TypeAlias = (
137137
)
138138
UIntDtypeArg: TypeAlias = (
139139
# Pandas nullable unsigned integer types and their string aliases
140-
pd.UInt8Dtype # noqa: PYI030
140+
pd.UInt8Dtype # noqa: PYI030,PYI055
141141
| pd.UInt16Dtype
142142
| pd.UInt32Dtype
143143
| pd.UInt64Dtype
@@ -166,7 +166,7 @@ UIntDtypeArg: TypeAlias = (
166166
)
167167
FloatDtypeArg: TypeAlias = (
168168
# Builtin float type and its string alias
169-
type[float] # noqa: PYI030
169+
type[float] # noqa: PYI030,PYI055
170170
| Literal["float"]
171171
# Pandas nullable float types and their string aliases
172172
| pd.Float32Dtype
@@ -197,7 +197,7 @@ FloatDtypeArg: TypeAlias = (
197197
)
198198
ComplexDtypeArg: TypeAlias = (
199199
# Builtin complex type and its string alias
200-
type[complex] # noqa: PYI030
200+
type[complex] # noqa: PYI030,PYI055
201201
| Literal["complex"]
202202
# Numpy complex types and their aliases
203203
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.csingle
@@ -326,7 +326,7 @@ TimestampDtypeArg: TypeAlias = Literal[
326326

327327
StrDtypeArg: TypeAlias = (
328328
# Builtin str type and its string alias
329-
type[str] # noqa: PYI030
329+
type[str] # noqa: PYI030,PYI055
330330
| Literal["str"]
331331
# Pandas nullable string type and its string alias
332332
| pd.StringDtype
@@ -340,7 +340,7 @@ StrDtypeArg: TypeAlias = (
340340
)
341341
BytesDtypeArg: TypeAlias = (
342342
# Builtin bytes type and its string alias
343-
type[bytes] # noqa: PYI030
343+
type[bytes] # noqa: PYI030,PYI055
344344
| Literal["bytes"]
345345
# Numpy bytes type and its string alias
346346
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.bytes_
@@ -353,7 +353,7 @@ CategoryDtypeArg: TypeAlias = CategoricalDtype | Literal["category"]
353353

354354
ObjectDtypeArg: TypeAlias = (
355355
# Builtin object type and its string alias
356-
type[object] # noqa: PYI030
356+
type[object] # noqa: PYI030,PYI055
357357
| Literal["object"]
358358
# Numpy object type and its string alias
359359
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.object_
@@ -483,6 +483,8 @@ ScalarT = TypeVar("ScalarT", bound=Scalar)
483483
np_ndarray_int64: TypeAlias = npt.NDArray[np.int64]
484484
np_ndarray_int: TypeAlias = npt.NDArray[np.signedinteger]
485485
np_ndarray_anyint: TypeAlias = npt.NDArray[np.integer]
486+
np_ndarray_float: TypeAlias = npt.NDArray[np.floating]
487+
np_ndarray_complex: TypeAlias = npt.NDArray[np.complexfloating]
486488
np_ndarray_bool: TypeAlias = npt.NDArray[np.bool_]
487489
np_ndarray_str: TypeAlias = npt.NDArray[np.str_]
488490

@@ -511,19 +513,12 @@ S1 = TypeVar(
511513
| float
512514
| complex
513515
| Dtype
514-
| Timestamp
515-
| Timedelta
516+
| datetime.datetime # includes pd.Timestamp
517+
| datetime.timedelta # includes pd.Timedelta
516518
| Period
517-
| Interval[int]
518-
| Interval[float]
519-
| Interval[Timestamp]
520-
| Interval[Timedelta]
519+
| Interval[int | float | Timestamp | Timedelta]
521520
| CategoricalDtype,
522521
)
523-
T1 = TypeVar(
524-
"T1", str, int, np.int64, np.uint64, np.float64, float, np.dtype[np.generic]
525-
)
526-
T2 = TypeVar("T2", str, int)
527522

528523
IndexingInt: TypeAlias = (
529524
int | np.int_ | np.integer | np.unsignedinteger | np.signedinteger | np.int8
@@ -601,14 +596,9 @@ ByT = TypeVar(
601596
| int
602597
| float
603598
| complex
604-
| Timestamp
605-
| Timedelta
606599
| Scalar
607600
| Period
608-
| Interval[int]
609-
| Interval[float]
610-
| Interval[Timestamp]
611-
| Interval[Timedelta]
601+
| Interval[int | float | Timestamp | Timedelta]
612602
| tuple,
613603
)
614604
# Use a distinct SeriesByT when using groupby with Series of known dtype.
@@ -622,13 +612,10 @@ SeriesByT = TypeVar(
622612
| int
623613
| float
624614
| complex
625-
| Timestamp
626-
| Timedelta
615+
| datetime.datetime
616+
| datetime.timedelta
627617
| Period
628-
| Interval[int]
629-
| Interval[float]
630-
| Interval[Timestamp]
631-
| Interval[Timedelta],
618+
| Interval[int | float | Timestamp | Timedelta],
632619
)
633620
GroupByObjectNonScalar: TypeAlias = (
634621
tuple

‎pandas-stubs/core/accessor.pyi

-3
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,5 @@
11
from typing import Any
22

3-
class DirNamesMixin:
4-
def __dir__(self): ...
5-
63
class PandasDelegate: ...
74

85
def delegate_names(

‎pandas-stubs/core/arraylike.pyi

+3-1
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,9 @@ class OpsMixin:
3131
def __rmul__(self, other: Any) -> Self: ...
3232
def __truediv__(self, other: Any) -> Self: ...
3333
def __rtruediv__(self, other: Any) -> Self: ...
34-
def __floordiv__(self, other: Any) -> Self: ...
34+
# __floordiv__ is handled by subclasses that specify only the valid values
35+
# that can be passed
36+
# def __floordiv__(self, other: Any) -> Self: ...
3537
def __rfloordiv__(self, other: Any) -> Self: ...
3638
def __mod__(self, other: Any) -> Self: ...
3739
def __rmod__(self, other: Any) -> Self: ...

‎pandas-stubs/core/arrays/categorical.pyi

+3-6
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,7 @@ import numpy as np
1212
from pandas import Series
1313
from pandas.core.accessor import PandasDelegate as PandasDelegate
1414
from pandas.core.arrays.base import ExtensionArray as ExtensionArray
15-
from pandas.core.base import (
16-
NoNewAttributesMixin as NoNewAttributesMixin,
17-
PandasObject as PandasObject,
18-
)
15+
from pandas.core.base import NoNewAttributesMixin as NoNewAttributesMixin
1916
from pandas.core.indexes.base import Index
2017

2118
from pandas._typing import (
@@ -33,7 +30,7 @@ from pandas.core.dtypes.dtypes import CategoricalDtype as CategoricalDtype
3330

3431
def contains(cat, key, container): ...
3532

36-
class Categorical(ExtensionArray, PandasObject):
33+
class Categorical(ExtensionArray):
3734
__array_priority__: int = ...
3835
def __init__(
3936
self,
@@ -185,7 +182,7 @@ class Categorical(ExtensionArray, PandasObject):
185182
def repeat(self, repeats, axis=...): ...
186183
def isin(self, values): ...
187184

188-
class CategoricalAccessor(PandasDelegate, PandasObject, NoNewAttributesMixin):
185+
class CategoricalAccessor(PandasDelegate, NoNewAttributesMixin):
189186
def __init__(self, data) -> None: ...
190187
@property
191188
def codes(self) -> Series[int]: ...

‎pandas-stubs/core/arrays/sparse/array.pyi

+1-2
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,10 @@ from pandas.core.arrays import (
33
ExtensionArray,
44
ExtensionOpsMixin,
55
)
6-
from pandas.core.base import PandasObject
76

87
from pandas._typing import TakeIndexer
98

10-
class SparseArray(PandasObject, ExtensionArray, ExtensionOpsMixin):
9+
class SparseArray(ExtensionArray, ExtensionOpsMixin):
1110
def __init__(
1211
self,
1312
data,

‎pandas-stubs/core/base.pyi

+11-14
Original file line numberDiff line numberDiff line change
@@ -1,43 +1,42 @@
11
from typing import (
22
Generic,
3+
Iterator,
34
Literal,
45
)
56

67
import numpy as np
78
from pandas import Index
8-
from pandas.core.accessor import DirNamesMixin
99
from pandas.core.arraylike import OpsMixin
1010
from pandas.core.arrays import ExtensionArray
1111
from pandas.core.arrays.categorical import Categorical
12+
from typing_extensions import Self
1213

1314
from pandas._typing import (
15+
S1,
1416
AxisIndex,
1517
NaPosition,
1618
NDFrameT,
1719
Scalar,
1820
npt,
1921
)
2022

21-
class PandasObject(DirNamesMixin):
22-
def __sizeof__(self) -> int: ...
23-
2423
class NoNewAttributesMixin:
2524
def __setattr__(self, key, value) -> None: ...
2625

2726
class SelectionMixin(Generic[NDFrameT]):
2827
def ndim(self) -> int: ...
2928
def __getitem__(self, key): ...
3029

31-
class IndexOpsMixin(OpsMixin):
30+
class IndexOpsMixin(OpsMixin, Generic[S1]):
3231
__array_priority__: int = ...
33-
def transpose(self, *args, **kwargs) -> IndexOpsMixin: ...
32+
def transpose(self, *args, **kwargs) -> Self: ...
3433
@property
35-
def T(self) -> IndexOpsMixin: ...
34+
def T(self) -> Self: ...
3635
@property
3736
def shape(self) -> tuple: ...
3837
@property
3938
def ndim(self) -> int: ...
40-
def item(self): ...
39+
def item(self) -> S1: ...
4140
@property
4241
def nbytes(self) -> int: ...
4342
@property
@@ -61,9 +60,9 @@ class IndexOpsMixin(OpsMixin):
6160
def argmin(
6261
self, axis: AxisIndex | None = ..., skipna: bool = ..., *args, **kwargs
6362
) -> np.int64: ...
64-
def tolist(self) -> list: ...
65-
def to_list(self) -> list: ...
66-
def __iter__(self): ...
63+
def tolist(self) -> list[S1]: ...
64+
def to_list(self) -> list[S1]: ...
65+
def __iter__(self) -> Iterator[S1]: ...
6766
@property
6867
def hasnans(self) -> bool: ...
6968
def value_counts(
@@ -87,6 +86,4 @@ class IndexOpsMixin(OpsMixin):
8786
def searchsorted(
8887
self, value, side: Literal["left", "right"] = ..., sorter=...
8988
) -> int | list[int]: ...
90-
def drop_duplicates(
91-
self, *, keep: NaPosition | Literal[False] = ...
92-
) -> IndexOpsMixin: ...
89+
def drop_duplicates(self, *, keep: NaPosition | Literal[False] = ...) -> Self: ...

‎pandas-stubs/core/computation/ops.pyi

-1
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@ class Term:
3434
def ndim(self) -> int: ...
3535

3636
class Constant(Term):
37-
def __init__(self, value, env, side=..., encoding=...) -> None: ...
3837
@property
3938
def name(self): ...
4039

‎pandas-stubs/core/computation/pytables.pyi

+1-1
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ class Term(ops.Term):
2828
def value(self, new_value) -> None: ...
2929

3030
class Constant(Term):
31-
def __init__(self, value, env: PyTablesScope, side=..., encoding=...) -> None: ...
31+
def __init__(self, name, env: PyTablesScope, side=..., encoding=...) -> None: ...
3232

3333
class BinOp(ops.BinOp):
3434
op: str

‎pandas-stubs/core/frame.pyi

+5-1
Original file line numberDiff line numberDiff line change
@@ -1482,7 +1482,7 @@ class DataFrame(NDFrame, OpsMixin):
14821482
@property
14831483
def at(self): ... # Not sure what to do with this yet; look at source
14841484
@property
1485-
def columns(self) -> Index: ...
1485+
def columns(self) -> Index[str]: ...
14861486
@columns.setter # setter needs to be right next to getter; otherwise mypy complains
14871487
def columns(
14881488
self, cols: AnyArrayLike | list[HashableT] | tuple[HashableT, ...]
@@ -2275,3 +2275,7 @@ class DataFrame(NDFrame, OpsMixin):
22752275
level: Level | None = ...,
22762276
drop_level: _bool = ...,
22772277
) -> DataFrame | Series: ...
2278+
# floordiv overload
2279+
def __floordiv__(
2280+
self, other: float | DataFrame | Series[int] | Series[float]
2281+
) -> Self: ...

‎pandas-stubs/core/generic.pyi

+1-2
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ from typing import (
1616

1717
import numpy as np
1818
from pandas import Index
19-
from pandas.core.base import PandasObject
2019
import pandas.core.indexing as indexing
2120
from pandas.core.series import Series
2221
import sqlalchemy.engine
@@ -54,7 +53,7 @@ from pandas.io.sql import SQLTable
5453
_bool = bool
5554
_str = str
5655

57-
class NDFrame(PandasObject, indexing.IndexingMixin):
56+
class NDFrame(indexing.IndexingMixin):
5857
__hash__: ClassVar[None] # type: ignore[assignment]
5958

6059
def set_flags(

0 commit comments

Comments
 (0)