Make IndexOpsMixin (and Index) generic

.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/_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_
@@ -511,19 +511,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 +594,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 +610,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/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, ...]

pandas-stubs/core/indexes/accessors.pyi

@@ -36,8 +36,7 @@ from pandas._typing import (
     np_ndarray_bool,
 )
 
-class Properties(PandasDelegate, NoNewAttributesMixin):
-    def __init__(self, data: Series, orig) -> None: ...
+class Properties(PandasDelegate, NoNewAttributesMixin): ...
 
 _DTFieldOpsReturnType = TypeVar("_DTFieldOpsReturnType", Series[int], Index[int])
 

pandas-stubs/core/indexes/base.pyi

@@ -53,6 +53,7 @@ from pandas._typing import (
     np_ndarray_anyint,
     np_ndarray_bool,
     np_ndarray_int64,
+    npt,
     type_t,
 )
 
@@ -66,9 +67,9 @@ class Index(IndexOpsMixin[S1]):
     @overload
     def __new__(  # type: ignore[misc]
         cls,
-        data: Iterable[int | np.integer],
+        data: Sequence[int | np.integer] | IndexOpsMixin[int] | npt.NDArray[np.integer],
         *,
-        dtype: Literal["int"] | type_t[int] | type_t[np.integer] = ...,
+        dtype: Literal["int"] | type_t[int | np.integer] = ...,
         copy: bool = ...,
         name=...,
         tupleize_cols: bool = ...,
@@ -79,7 +80,7 @@ class Index(IndexOpsMixin[S1]):
         cls,
         data: Iterable,
         *,
-        dtype: Literal["int"] | type_t[int] | type_t[np.integer],
+        dtype: Literal["int"] | type_t[int | np.integer],
         copy: bool = ...,
         name=...,
         tupleize_cols: bool = ...,
@@ -88,12 +89,11 @@ class Index(IndexOpsMixin[S1]):
     @overload
     def __new__(  # type: ignore[misc]
         cls,
-        data: Iterable[float | np.float32 | np.float64],
+        data: Sequence[float | np.floating]
+        | IndexOpsMixin[float]
+        | npt.NDArray[np.floating],
         *,
-        dtype: Literal["float"]
-        | type_t[float]
-        | type_t[np.float32]
-        | type_t[np.float64] = ...,
+        dtype: Literal["float"] | type_t[float | np.floating] = ...,
         copy: bool = ...,
         name=...,
         tupleize_cols: bool = ...,
@@ -104,10 +104,7 @@ class Index(IndexOpsMixin[S1]):
         cls,
         data: Iterable,
         *,
-        dtype: Literal["float"]
-        | type_t[float]
-        | type_t[np.float32]
-        | type_t[np.float64],
+        dtype: Literal["float"] | type_t[float | np.floating],
         copy: bool = ...,
         name=...,
         tupleize_cols: bool = ...,
@@ -116,9 +113,11 @@ class Index(IndexOpsMixin[S1]):
     @overload
     def __new__(  # type: ignore[misc]
         cls,
-        data: Iterable[complex],
+        data: Sequence[complex | np.complexfloating]
+        | IndexOpsMixin[complex]
+        | npt.NDArray[np.complexfloating],
         *,
-        dtype: Literal["complex"] | type_t[complex] = ...,
+        dtype: Literal["complex"] | type_t[complex | np.complexfloating] = ...,
         copy: bool = ...,
         name=...,
         tupleize_cols: bool = ...,
@@ -129,7 +128,7 @@ class Index(IndexOpsMixin[S1]):
         cls,
         data: Iterable,
         *,
-        dtype: Literal["complex"] | type_t[complex],
+        dtype: Literal["complex"] | type_t[complex | np.complexfloating],
         copy: bool = ...,
         name=...,
         tupleize_cols: bool = ...,
@@ -139,7 +138,7 @@ class Index(IndexOpsMixin[S1]):
     @overload
     def __new__(  # type: ignore[misc]
         cls,
-        data: Iterable[np.datetime64 | datetime],
+        data: Sequence[np.datetime64 | datetime] | IndexOpsMixin[datetime],
         *,
         dtype: TimestampDtypeArg = ...,
         copy: bool = ...,
@@ -161,7 +160,7 @@ class Index(IndexOpsMixin[S1]):
     @overload
     def __new__(  # type: ignore[misc]
         cls,
-        data: Iterable[Period],
+        data: Sequence[Period] | IndexOpsMixin[Period],
         *,
         dtype: PeriodDtype = ...,
         copy: bool = ...,
@@ -183,7 +182,7 @@ class Index(IndexOpsMixin[S1]):
     @overload
     def __new__(  # type: ignore[misc]
         cls,
-        data: Iterable[np.timedelta64 | timedelta],
+        data: Sequence[np.timedelta64 | timedelta] | IndexOpsMixin[timedelta],
         *,
         dtype: TimedeltaDtypeArg = ...,
         copy: bool = ...,
@@ -205,7 +204,8 @@ class Index(IndexOpsMixin[S1]):
     @overload
     def __new__(  # type: ignore[misc]
         cls,
-        data: Iterable[Interval[_OrderableT]],
+        data: Sequence[Interval[_OrderableT]]  # type: ignore[type-var]
+        | IndexOpsMixin[Interval[_OrderableT]],
         *,
         dtype: Literal["Interval"] = ...,
         copy: bool = ...,
@@ -228,7 +228,7 @@ class Index(IndexOpsMixin[S1]):
     @overload
     def __new__(
         cls,
-        data: Iterable[S1] = ...,
+        data: Iterable[S1] | IndexOpsMixin[S1] = ...,
         *,
         dtype: type[S1] = ...,
         copy: bool = ...,

pandas-stubs/core/indexes/multi.pyi

@@ -15,7 +15,6 @@ from pandas.core.indexes.base import Index
 from typing_extensions import Self
 
 from pandas._typing import (
-    T1,
     Dtype,
     DtypeArg,
     HashableT,
@@ -156,7 +155,7 @@ class MultiIndex(Index[Any]):
     def union(self, other, sort=...): ...
     def intersection(self, other, sort: bool = ...): ...
     def difference(self, other, sort=...): ...
-    def astype(self, dtype: DtypeArg | T1, copy: bool = ...) -> Self: ...
+    def astype(self, dtype: DtypeArg, copy: bool = ...) -> Self: ...
     def insert(self, loc, item): ...
     def delete(self, loc): ...
     def isin(self, values, level=...) -> np_ndarray_bool: ...

pandas-stubs/core/indexes/timedeltas.pyi

@@ -32,18 +32,6 @@ from pandas._typing import (
 )
 
 class TimedeltaIndex(DatetimeTimedeltaMixin[Timedelta], TimedeltaIndexProperties):
-    def __init__(
-        self,
-        data: AnyArrayLike
-        | list[str]
-        | Sequence[dt.timedelta | Timedelta | np.timedelta64 | float] = ...,
-        unit: Literal["D", "h", "m", "s", "ms", "us", "ns"] = ...,
-        freq: str | BaseOffset = ...,
-        closed: object = ...,
-        dtype: Literal["<m8[ns]"] = ...,
-        copy: bool = ...,
-        name: str = ...,
-    ) -> None: ...
     def __new__(
         cls,
         data: AnyArrayLike

pyproject.toml

@@ -65,6 +65,7 @@ jinja2 = ">=3.1"
 scipy = ">=1.9.1"
 SQLAlchemy = ">=2.0.12"
 types-python-dateutil = ">=2.8.19"
+numexpr = "<2.8.5"  # https://github.com/pandas-dev/pandas/issues/54449
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]
@@ -195,7 +196,6 @@ reportPrivateUsage = false
 # enable optional checks
 reportMissingModuleSource = true
 useLibraryCodeForTypes = false
-defineConstant = { MYPY_CHECKING = false }
 
 [tool.codespell]
 ignore-words-list = "indext, mose, sav, ser"

tests/test_frame.py

@@ -1761,8 +1761,7 @@ def test_getmultiindex_columns() -> None:
         [(i, s) for i in [1] for s in df.columns.get_level_values(1)]
     ]
     res4: pd.DataFrame = df[[df.columns[0]]]
-    column: Scalar = df.columns[0]
-    check(assert_type(df[column], pd.Series), pd.Series)
+    check(assert_type(df[df.columns[0]], pd.Series), pd.Series)
     check(assert_type(df[li[0]], pd.Series), pd.Series)
 
 
@@ -2592,7 +2591,7 @@ def test_in_columns() -> None:
     # GH 532 (PR)
     df = pd.DataFrame(np.random.random((3, 4)), columns=["cat", "dog", "rat", "pig"])
     cols = [c for c in df.columns if "at" in c]
-    check(assert_type(cols, list), list, str)
+    check(assert_type(cols, "list[str]"), list, str)
     check(assert_type(df.loc[:, cols], pd.DataFrame), pd.DataFrame)
     check(assert_type(df[cols], pd.DataFrame), pd.DataFrame)
     check(assert_type(df.groupby(by=cols).sum(), pd.DataFrame), pd.DataFrame)

tests/test_indexes.py

@@ -2,7 +2,6 @@
 
 import datetime as dt
 from typing import (
-    TYPE_CHECKING,
     Tuple,
     Union,
 )
@@ -22,9 +21,6 @@
     check,
 )
 
-if TYPE_CHECKING:
-    MYPY_CHECKING: bool = True
-
 
 def test_index_unique() -> None:
     df = pd.DataFrame({"x": [1, 2, 3, 4]}, index=pd.Index([1, 2, 3, 2]))
@@ -77,9 +73,9 @@ def test_column_getitem() -> None:
     # https://github.com/microsoft/python-type-stubs/issues/199#issuecomment-1132806594
     df = pd.DataFrame([[1, 2, 3]], columns=["a", "b", "c"])
 
-    columns: pd.Index[str] = df.columns
-    check(assert_type(columns[0], str), str)
-    check(assert_type(df[columns[0]], pd.Series), pd.Series, np.int64)
+    column = df.columns[0]
+    check(assert_type(column, str), str)
+    check(assert_type(df[column], pd.Series), pd.Series, np.int64)
 
 
 def test_column_contains() -> None:
@@ -97,7 +93,7 @@ def test_column_sequence() -> None:
     df = pd.DataFrame([1, 2, 3])
     col_list = list(df.columns)
     check(
-        assert_type(col_list, list),
+        assert_type(col_list, "list[str]"),
         list,
         int,
     )