Type RangeIndex (pandas-dev#1453)

* type rangeindex

* add tests

* get_indexer

* argsort

* type join

* update roof

* consistency

* type default

Author: MarcoGorelli

pandas-stubs/core/indexes/base.pyi

@@ -66,6 +66,7 @@ from pandas._typing import (
     S2_NSDT,
     T_COMPLEX,
     AnyAll,
+    AnyArrayLike,
     ArrayLike,
     AxesData,
     CategoryDtypeArg,
@@ -78,6 +79,7 @@ from pandas._typing import (
     GenericT_co,
     HashableT,
     IgnoreRaise,
+    JoinHow,
     Just,
     Label,
     Level,
@@ -418,8 +420,12 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
     ) -> Self: ...
     def get_loc(self, key: Label) -> int | slice | np_1darray[np.bool]: ...
     def get_indexer(
-        self, target, method: ReindexMethod | None = ..., limit=..., tolerance=...
-    ): ...
+        self,
+        target: Index,
+        method: ReindexMethod | None = None,
+        limit: int | None = None,
+        tolerance: Scalar | AnyArrayLike | Sequence[Scalar] | None = None,
+    ) -> np_1darray[np.intp]: ...
     def reindex(
         self,
         target,
@@ -428,15 +434,26 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
         limit=...,
         tolerance=...,
     ): ...
+    @overload
     def join(
         self,
-        other,
+        other: Index,
         *,
-        how: _str = ...,
-        level=...,
-        return_indexers: bool = ...,
-        sort: bool = ...,
-    ): ...
+        how: JoinHow = "left",
+        level: Level | None = None,
+        return_indexers: Literal[True],
+        sort: bool = False,
+    ) -> tuple[Index, np_1darray[np.intp] | None, np_1darray[np.intp] | None]: ...
+    @overload
+    def join(
+        self,
+        other: Index,
+        *,
+        how: JoinHow = "left",
+        level: Level | None = None,
+        return_indexers: Literal[False] = False,
+        sort: bool = False,
+    ) -> Index: ...
     @property
     def values(self) -> np_1darray: ...
     def memory_usage(self, deep: bool = False): ...
@@ -474,7 +491,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
     ): ...
     @final
     def sort(self, *args: Any, **kwargs: Any) -> None: ...
-    def argsort(self, *args: Any, **kwargs: Any): ...
+    def argsort(self, *args: Any, **kwargs: Any) -> np_1darray[np.intp]: ...
     def get_indexer_non_unique(self, target): ...
     @final
     def get_indexer_for(self, target, **kwargs: Any): ...

pandas-stubs/core/indexes/interval.pyi

@@ -226,7 +226,7 @@ class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin):
         method: FillnaOptions | Literal["nearest"] | None = ...,
         limit: int | None = ...,
         tolerance=...,
-    ) -> npt.NDArray[np.intp]: ...
+    ) -> np_1darray[np.intp]: ...
     def get_indexer_non_unique(
         self, target: Index
     ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: ...

pandas-stubs/core/indexes/range.pyi

@@ -16,6 +16,7 @@ from pandas.core.indexes.base import (
 from typing_extensions import Self
 
 from pandas._typing import (
+    Dtype,
     HashableT,
     MaskType,
     np_1darray,
@@ -25,16 +26,17 @@ from pandas._typing import (
 class RangeIndex(_IndexSubclassBase[int, np.int64]):
     def __new__(
         cls,
-        start: int | RangeIndex | range = ...,
-        stop: int = ...,
-        step: int = ...,
-        dtype=...,
-        copy: bool = ...,
-        name: Hashable = ...,
+        start: int | RangeIndex | range | None = None,
+        stop: int | None = None,
+        step: int | None = None,
+        dtype: Dtype | None = None,
+        copy: bool = False,
+        name: Hashable | None = None,
     ): ...
     @classmethod
-    def from_range(cls, data, name: Hashable = ..., dtype=...) -> Self: ...
-    def __reduce__(self): ...
+    def from_range(
+        cls, data: range, name: Hashable | None = None, dtype: Dtype | None = None
+    ) -> Self: ...
     @property
     def start(self) -> int: ...
     @property
@@ -55,29 +57,15 @@ class RangeIndex(_IndexSubclassBase[int, np.int64]):
     @property
     def has_duplicates(self) -> bool: ...
     def __contains__(self, key: int | np.integer) -> bool: ...
-    @final
-    def get_indexer(self, target, method=..., limit=..., tolerance=...): ...
-    def tolist(self): ...
-    def min(self, axis=..., skipna: bool = ..., *args: Any, **kwargs: Any): ...
-    def max(self, axis=..., skipna: bool = ..., *args: Any, **kwargs: Any): ...
-    def argsort(self, *args: Any, **kwargs: Any): ...
     def factorize(
         self, sort: bool = False, use_na_sentinel: bool = True
     ) -> tuple[np_1darray[np.intp], RangeIndex]: ...
-    def equals(self, other): ...
-    @final
-    def join(
-        self,
-        other,
-        *,
-        how: str = ...,
-        level=...,
-        return_indexers: bool = ...,
-        sort: bool = ...,
-    ): ...
     @property
     def size(self) -> int: ...
-    def __floordiv__(self, other): ...
+    # Base class returns `Self`, but for `RangeIndex` that's not true.
+    def __floordiv__(  # type: ignore[override]
+        self, other: float | Sequence[float] | Index[int] | Index[float]
+    ) -> Index[int]: ...
     def all(self, *args: Any, **kwargs: Any) -> bool: ...
     def any(self, *args: Any, **kwargs: Any) -> bool: ...
     @final

pyproject.toml

@@ -204,10 +204,6 @@ ignore = [
   "PYI042",  # https://docs.astral.sh/ruff/rules/snake-case-type-alias/
   "ERA001", "PLR0402", "PLC0105"
 ]
-"*range.pyi" = [
-  # TODO: remove when pandas-dev/pandas-stubs#1442 is resolved
-  "ANN001", "ANN201", "ANN204", "ANN206",
-]
 "*category.pyi" = [
   # TODO: remove when pandas-dev/pandas-stubs#1443 is resolved
   "ANN001", "ANN201", "ANN204", "ANN206",

tests/indexes/test_rangeindex.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+from typing_extensions import (
+    assert_type,
+)
+
+from tests import (
+    check,
+    np_1darray,
+)
+
+
+def test_rangeindex_floordiv() -> None:
+    ri = pd.RangeIndex(3)
+    check(
+        assert_type(ri // 2, "pd.Index[int]"),
+        pd.Index,
+    )
+
+
+def test_rangeindex_min_max() -> None:
+    ri = pd.RangeIndex(3)
+    check(
+        assert_type(ri.min(), int),
+        int,
+    )
+    check(
+        assert_type(ri.max(axis=0), int),
+        int,
+    )
+
+
+def test_rangeindex_equals() -> None:
+    ri = pd.RangeIndex(3)
+    check(
+        assert_type(ri.equals(ri), bool),
+        bool,
+    )
+
+
+def test_rangeindex_tolist() -> None:
+    ri = pd.RangeIndex.from_range(range(3))
+    check(
+        assert_type(ri.tolist(), list[int]),
+        list[int],
+    )
+
+
+def test_rangeindex_get_indexer() -> None:
+    ri = pd.RangeIndex.from_range(range(3))
+    check(
+        assert_type(ri.get_indexer(ri), np_1darray[np.intp]),
+        np_1darray[np.intp],
+    )
+
+
+def test_rangeindex_argsort() -> None:
+    ri = pd.RangeIndex.from_range(range(3))
+    check(
+        assert_type(ri.argsort(), np_1darray[np.intp]),
+        np_1darray[np.intp],
+    )
+
+
+def test_rangeindex_join() -> None:
+    ri = pd.RangeIndex.from_range(range(3))
+    check(
+        assert_type(ri.join(ri), pd.Index),
+        pd.Index,
+    )
+    check(
+        assert_type(
+            ri.join(ri, return_indexers=True),
+            tuple[pd.Index, np_1darray[np.intp] | None, np_1darray[np.intp] | None],
+        ),
+        tuple[pd.Index, np_1darray[np.intp] | None, np_1darray[np.intp] | None],
+    )