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 <irv@princeton.com>

Author: twoertwein

.pre-commit-config.yaml

@@ -11,7 +11,7 @@ repos:
     hooks:
     -   id: isort
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.0.282
+    rev: v0.0.283
     hooks:
     -   id: ruff
         args: [

pandas-stubs/_libs/tslibs/timedeltas.pyi

@@ -11,14 +11,11 @@ import numpy as np
 import pandas as pd
 from pandas import (
     DatetimeIndex,
+    Index,
     PeriodIndex,
     Series,
     TimedeltaIndex,
 )
-from pandas.core.indexes.base import (
-    _FloatIndexType,
-    _IntIndexType,
-)
 from pandas.core.series import (
     TimedeltaSeries,
     TimestampSeries,
@@ -245,7 +242,7 @@ class Timedelta(timedelta):
     @overload
     def __mul__(self, other: Series[float]) -> TimedeltaSeries: ...
     @overload
-    def __mul__(self, other: _IntIndexType | _FloatIndexType) -> TimedeltaIndex: ...
+    def __mul__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
     @overload
     def __rmul__(self, other: float) -> Timedelta: ...
     @overload
@@ -254,8 +251,9 @@ class Timedelta(timedelta):
     def __rmul__(self, other: Series[int]) -> TimedeltaSeries: ...
     @overload
     def __rmul__(self, other: Series[float]) -> TimedeltaSeries: ...
+    # maybe related to https://github.com/python/mypy/issues/10755
     @overload
-    def __rmul__(self, other: _IntIndexType | _FloatIndexType) -> TimedeltaIndex: ...
+    def __rmul__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...  # type: ignore[misc]
     # Override due to more types supported than dt.timedelta
     # error: Signature of "__floordiv__" incompatible with supertype "timedelta"
     @overload  # type: ignore[override]
@@ -271,9 +269,7 @@ class Timedelta(timedelta):
         self, other: npt.NDArray[np.timedelta64]
     ) -> npt.NDArray[np.int_]: ...
     @overload
-    def __floordiv__(
-        self, other: _IntIndexType | _FloatIndexType
-    ) -> TimedeltaIndex: ...
+    def __floordiv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
     @overload
     def __floordiv__(self, other: Series[int]) -> TimedeltaSeries: ...
     @overload
@@ -306,7 +302,7 @@ class Timedelta(timedelta):
     @overload
     def __truediv__(self, other: Series[float]) -> TimedeltaSeries: ...
     @overload
-    def __truediv__(self, other: _IntIndexType | _FloatIndexType) -> TimedeltaIndex: ...
+    def __truediv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
     def __rtruediv__(self, other: timedelta | Timedelta | NaTType) -> float: ...
     # Override due to more types supported than dt.timedelta
     @overload
@@ -338,7 +334,7 @@ class Timedelta(timedelta):
     @overload
     def __mod__(self, other: Series[int] | Series[float]) -> TimedeltaSeries: ...
     @overload
-    def __mod__(self, other: _IntIndexType | _FloatIndexType) -> TimedeltaIndex: ...
+    def __mod__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
     @overload
     def __mod__(
         self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]

pandas-stubs/_typing.pyi

@@ -91,7 +91,7 @@ DtypeBackend: TypeAlias = Literal["pyarrow", "numpy_nullable"]
 
 BooleanDtypeArg: TypeAlias = (
     # Builtin bool type and its string alias
-    type[bool]  # noqa: PYI030
+    type[bool]  # noqa: PYI030,PYI055
     | Literal["bool"]
     # Pandas nullable boolean type and its string alias
     | pd.BooleanDtype
@@ -105,7 +105,7 @@ BooleanDtypeArg: TypeAlias = (
 )
 IntDtypeArg: TypeAlias = (
     # Builtin integer type and its string alias
-    type[int]  # noqa: PYI030
+    type[int]  # noqa: PYI030,PYI055
     | Literal["int"]
     # Pandas nullable integer types and their string aliases
     | pd.Int8Dtype
@@ -137,7 +137,7 @@ IntDtypeArg: TypeAlias = (
 )
 UIntDtypeArg: TypeAlias = (
     # Pandas nullable unsigned integer types and their string aliases
-    pd.UInt8Dtype  # noqa: PYI030
+    pd.UInt8Dtype  # noqa: PYI030,PYI055
     | pd.UInt16Dtype
     | pd.UInt32Dtype
     | pd.UInt64Dtype
@@ -166,7 +166,7 @@ UIntDtypeArg: TypeAlias = (
 )
 FloatDtypeArg: TypeAlias = (
     # Builtin float type and its string alias
-    type[float]  # noqa: PYI030
+    type[float]  # noqa: PYI030,PYI055
     | Literal["float"]
     # Pandas nullable float types and their string aliases
     | pd.Float32Dtype
@@ -197,7 +197,7 @@ FloatDtypeArg: TypeAlias = (
 )
 ComplexDtypeArg: TypeAlias = (
     # Builtin complex type and its string alias
-    type[complex]  # noqa: PYI030
+    type[complex]  # noqa: PYI030,PYI055
     | Literal["complex"]
     # Numpy complex types and their aliases
     # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.csingle
@@ -326,7 +326,7 @@ TimestampDtypeArg: TypeAlias = Literal[
 
 StrDtypeArg: TypeAlias = (
     # Builtin str type and its string alias
-    type[str]  # noqa: PYI030
+    type[str]  # noqa: PYI030,PYI055
     | Literal["str"]
     # Pandas nullable string type and its string alias
     | pd.StringDtype
@@ -340,7 +340,7 @@ StrDtypeArg: TypeAlias = (
 )
 BytesDtypeArg: TypeAlias = (
     # Builtin bytes type and its string alias
-    type[bytes]  # noqa: PYI030
+    type[bytes]  # noqa: PYI030,PYI055
     | Literal["bytes"]
     # Numpy bytes type and its string alias
     # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.bytes_
@@ -353,7 +353,7 @@ CategoryDtypeArg: TypeAlias = CategoricalDtype | Literal["category"]
 
 ObjectDtypeArg: TypeAlias = (
     # Builtin object type and its string alias
-    type[object]  # noqa: PYI030
+    type[object]  # noqa: PYI030,PYI055
     | Literal["object"]
     # Numpy object type and its string alias
     # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.object_
@@ -483,6 +483,8 @@ ScalarT = TypeVar("ScalarT", bound=Scalar)
 np_ndarray_int64: TypeAlias = npt.NDArray[np.int64]
 np_ndarray_int: TypeAlias = npt.NDArray[np.signedinteger]
 np_ndarray_anyint: TypeAlias = npt.NDArray[np.integer]
+np_ndarray_float: TypeAlias = npt.NDArray[np.floating]
+np_ndarray_complex: TypeAlias = npt.NDArray[np.complexfloating]
 np_ndarray_bool: TypeAlias = npt.NDArray[np.bool_]
 np_ndarray_str: TypeAlias = npt.NDArray[np.str_]
 
@@ -511,19 +513,12 @@ S1 = TypeVar(
     | float
     | complex
     | Dtype
-    | Timestamp
-    | Timedelta
+    | datetime.datetime  # includes pd.Timestamp
+    | datetime.timedelta  # includes pd.Timedelta
     | Period
-    | Interval[int]
-    | Interval[float]
-    | Interval[Timestamp]
-    | Interval[Timedelta]
+    | Interval[int | float | Timestamp | Timedelta]
     | CategoricalDtype,
 )
-T1 = TypeVar(
-    "T1", str, int, np.int64, np.uint64, np.float64, float, np.dtype[np.generic]
-)
-T2 = TypeVar("T2", str, int)
 
 IndexingInt: TypeAlias = (
     int | np.int_ | np.integer | np.unsignedinteger | np.signedinteger | np.int8
@@ -601,14 +596,9 @@ ByT = TypeVar(
     | int
     | float
     | complex
-    | Timestamp
-    | Timedelta
     | Scalar
     | Period
-    | Interval[int]
-    | Interval[float]
-    | Interval[Timestamp]
-    | Interval[Timedelta]
+    | Interval[int | float | Timestamp | Timedelta]
     | tuple,
 )
 # Use a distinct SeriesByT when using groupby with Series of known dtype.
@@ -622,13 +612,10 @@ SeriesByT = TypeVar(
     | int
     | float
     | complex
-    | Timestamp
-    | Timedelta
+    | datetime.datetime
+    | datetime.timedelta
     | Period
-    | Interval[int]
-    | Interval[float]
-    | Interval[Timestamp]
-    | Interval[Timedelta],
+    | Interval[int | float | Timestamp | Timedelta],
 )
 GroupByObjectNonScalar: TypeAlias = (
     tuple

pandas-stubs/core/accessor.pyi

@@ -1,8 +1,5 @@
 from typing import Any
 
-class DirNamesMixin:
-    def __dir__(self): ...
-
 class PandasDelegate: ...
 
 def delegate_names(

pandas-stubs/core/arraylike.pyi

@@ -31,7 +31,9 @@ class OpsMixin:
     def __rmul__(self, other: Any) -> Self: ...
     def __truediv__(self, other: Any) -> Self: ...
     def __rtruediv__(self, other: Any) -> Self: ...
-    def __floordiv__(self, other: Any) -> Self: ...
+    # __floordiv__ is handled by subclasses that specify only the valid values
+    # that can be passed
+    #    def __floordiv__(self, other: Any) -> Self: ...
     def __rfloordiv__(self, other: Any) -> Self: ...
     def __mod__(self, other: Any) -> Self: ...
     def __rmod__(self, other: Any) -> Self: ...

pandas-stubs/core/arrays/categorical.pyi

@@ -12,10 +12,7 @@ import numpy as np
 from pandas import Series
 from pandas.core.accessor import PandasDelegate as PandasDelegate
 from pandas.core.arrays.base import ExtensionArray as ExtensionArray
-from pandas.core.base import (
-    NoNewAttributesMixin as NoNewAttributesMixin,
-    PandasObject as PandasObject,
-)
+from pandas.core.base import NoNewAttributesMixin as NoNewAttributesMixin
 from pandas.core.indexes.base import Index
 
 from pandas._typing import (
@@ -33,7 +30,7 @@ from pandas.core.dtypes.dtypes import CategoricalDtype as CategoricalDtype
 
 def contains(cat, key, container): ...
 
-class Categorical(ExtensionArray, PandasObject):
+class Categorical(ExtensionArray):
     __array_priority__: int = ...
     def __init__(
         self,
@@ -185,7 +182,7 @@ class Categorical(ExtensionArray, PandasObject):
     def repeat(self, repeats, axis=...): ...
     def isin(self, values): ...
 
-class CategoricalAccessor(PandasDelegate, PandasObject, NoNewAttributesMixin):
+class CategoricalAccessor(PandasDelegate, NoNewAttributesMixin):
     def __init__(self, data) -> None: ...
     @property
     def codes(self) -> Series[int]: ...

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

@@ -3,11 +3,10 @@ from pandas.core.arrays import (
     ExtensionArray,
     ExtensionOpsMixin,
 )
-from pandas.core.base import PandasObject
 
 from pandas._typing import TakeIndexer
 
-class SparseArray(PandasObject, ExtensionArray, ExtensionOpsMixin):
+class SparseArray(ExtensionArray, ExtensionOpsMixin):
     def __init__(
         self,
         data,

pandas-stubs/core/base.pyi

@@ -1,43 +1,42 @@
 from typing import (
     Generic,
+    Iterator,
     Literal,
 )
 
 import numpy as np
 from pandas import Index
-from pandas.core.accessor import DirNamesMixin
 from pandas.core.arraylike import OpsMixin
 from pandas.core.arrays import ExtensionArray
 from pandas.core.arrays.categorical import Categorical
+from typing_extensions import Self
 
 from pandas._typing import (
+    S1,
     AxisIndex,
     NaPosition,
     NDFrameT,
     Scalar,
     npt,
 )
 
-class PandasObject(DirNamesMixin):
-    def __sizeof__(self) -> int: ...
-
 class NoNewAttributesMixin:
     def __setattr__(self, key, value) -> None: ...
 
 class SelectionMixin(Generic[NDFrameT]):
     def ndim(self) -> int: ...
     def __getitem__(self, key): ...
 
-class IndexOpsMixin(OpsMixin):
+class IndexOpsMixin(OpsMixin, Generic[S1]):
     __array_priority__: int = ...
-    def transpose(self, *args, **kwargs) -> IndexOpsMixin: ...
+    def transpose(self, *args, **kwargs) -> Self: ...
     @property
-    def T(self) -> IndexOpsMixin: ...
+    def T(self) -> Self: ...
     @property
     def shape(self) -> tuple: ...
     @property
     def ndim(self) -> int: ...
-    def item(self): ...
+    def item(self) -> S1: ...
     @property
     def nbytes(self) -> int: ...
     @property
@@ -61,9 +60,9 @@ class IndexOpsMixin(OpsMixin):
     def argmin(
         self, axis: AxisIndex | None = ..., skipna: bool = ..., *args, **kwargs
     ) -> np.int64: ...
-    def tolist(self) -> list: ...
-    def to_list(self) -> list: ...
-    def __iter__(self): ...
+    def tolist(self) -> list[S1]: ...
+    def to_list(self) -> list[S1]: ...
+    def __iter__(self) -> Iterator[S1]: ...
     @property
     def hasnans(self) -> bool: ...
     def value_counts(
@@ -87,6 +86,4 @@ class IndexOpsMixin(OpsMixin):
     def searchsorted(
         self, value, side: Literal["left", "right"] = ..., sorter=...
     ) -> int | list[int]: ...
-    def drop_duplicates(
-        self, *, keep: NaPosition | Literal[False] = ...
-    ) -> IndexOpsMixin: ...
+    def drop_duplicates(self, *, keep: NaPosition | Literal[False] = ...) -> Self: ...

pandas-stubs/core/computation/ops.pyi

@@ -34,7 +34,6 @@ class Term:
     def ndim(self) -> int: ...
 
 class Constant(Term):
-    def __init__(self, value, env, side=..., encoding=...) -> None: ...
     @property
     def name(self): ...
 

pandas-stubs/core/computation/pytables.pyi

@@ -28,7 +28,7 @@ class Term(ops.Term):
     def value(self, new_value) -> None: ...
 
 class Constant(Term):
-    def __init__(self, value, env: PyTablesScope, side=..., encoding=...) -> None: ...
+    def __init__(self, name, env: PyTablesScope, side=..., encoding=...) -> None: ...
 
 class BinOp(ops.BinOp):
     op: str

pandas-stubs/core/frame.pyi

@@ -1482,7 +1482,7 @@ class DataFrame(NDFrame, OpsMixin):
     @property
     def at(self): ...  # Not sure what to do with this yet; look at source
     @property
-    def columns(self) -> Index: ...
+    def columns(self) -> Index[str]: ...
     @columns.setter  # setter needs to be right next to getter; otherwise mypy complains
     def columns(
         self, cols: AnyArrayLike | list[HashableT] | tuple[HashableT, ...]
@@ -2275,3 +2275,7 @@ class DataFrame(NDFrame, OpsMixin):
         level: Level | None = ...,
         drop_level: _bool = ...,
     ) -> DataFrame | Series: ...
+    # floordiv overload
+    def __floordiv__(
+        self, other: float | DataFrame | Series[int] | Series[float]
+    ) -> Self: ...

pandas-stubs/core/generic.pyi

@@ -16,7 +16,6 @@ from typing import (
 
 import numpy as np
 from pandas import Index
-from pandas.core.base import PandasObject
 import pandas.core.indexing as indexing
 from pandas.core.series import Series
 import sqlalchemy.engine
@@ -54,7 +53,7 @@ from pandas.io.sql import SQLTable
 _bool = bool
 _str = str
 
-class NDFrame(PandasObject, indexing.IndexingMixin):
+class NDFrame(indexing.IndexingMixin):
     __hash__: ClassVar[None]  # type: ignore[assignment]
 
     def set_flags(