diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 40129638a..25c975eff 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/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): ... diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index 628d588ec..9007c345e 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/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]]: ... diff --git a/pandas-stubs/core/indexes/range.pyi b/pandas-stubs/core/indexes/range.pyi index 781123313..c55e7f844 100644 --- a/pandas-stubs/core/indexes/range.pyi +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 91591a5e2..3876f5c77 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/indexes/test_rangeindex.py b/tests/indexes/test_rangeindex.py new file mode 100644 index 000000000..1b30e83ae --- /dev/null +++ b/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], + )