Type Index methods: is_, ravel, take, repeat, to_flat_index, to_frame, set_names, rename, remove undocumented __copy__, __deepcopy__, __array_wrap__ (#1486)

* Type Index methods: `is_`, `ravel`, `take`, `repeat`, `to_flat_index`,
`to_frame`, `set_names`, `rename`, remove undocumented `__copy__`,
`__deepcopy__`, `__array_wrap__`

* fix typing

* only allow for `Mapping` `names` in `set_names` for `MultiIndex`

* Apply suggestions from code review

Co-authored-by: Yi-Fan Wang <[email protected]>

* lint

* Update pandas-stubs/core/indexes/base.pyi

Co-authored-by: Yi-Fan Wang <[email protected]>

* fix no_default usage

* fix the extra `pd.Index([(1,)])` case

---------

Co-authored-by: Yi-Fan Wang <[email protected]>

Author: MarcoGorelli

pandas-stubs/core/indexes/base.pyi

@@ -3,7 +3,6 @@ from collections.abc import (
     Callable,
     Hashable,
     Iterable,
-    MutableMapping,
     Sequence,
 )
 from datetime import (
@@ -82,6 +81,7 @@ from pandas._typing import (
     AnyArrayLikeInt,
     ArrayLike,
     AxesData,
+    Axis,
     CategoryDtypeArg,
     DropKeep,
     Dtype,
@@ -104,6 +104,7 @@ from pandas._typing import (
     SequenceNotStr,
     SliceType,
     SupportsDType,
+    TakeIndexer,
     TimedeltaDtypeArg,
     TimestampDtypeArg,
     np_1darray,
@@ -350,40 +351,36 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
         Index,
     ]: ...
     @final
-    def is_(self, other) -> bool: ...
+    def is_(self, other: object) -> bool: ...
     def __len__(self) -> int: ...
     def __array__(
         self, dtype: _str | np.dtype = ..., copy: bool | None = ...
     ) -> np_1darray: ...
-    def __array_wrap__(self, result, context=...): ...
     @property
     def dtype(self) -> DtypeObj: ...
     @final
-    def ravel(self, order: _str = ...): ...
+    def ravel(self, order: _str = "C") -> Self: ...
     def view(self, cls=...): ...
     def astype(self, dtype: DtypeArg, copy: bool = True) -> Index: ...
     def take(
         self,
-        indices,
-        axis: int = 0,
+        indices: TakeIndexer,
+        axis: Axis = 0,
         allow_fill: bool = True,
         fill_value: Scalar | None = None,
         **kwargs: Any,
-    ): ...
-    def repeat(self, repeats, axis=...): ...
+    ) -> Self: ...
+    def repeat(
+        self, repeats: int | AnyArrayLikeInt | Sequence[int], axis: None = None
+    ) -> Self: ...
     def copy(self, name: Hashable = ..., deep: bool = False) -> Self: ...
-    @final
-    def __copy__(self, **kwargs: Any): ...
-    @final
-    def __deepcopy__(self, memo: MutableMapping[int, Any] | None = None) -> Self: ...
     def format(
         self, name: bool = ..., formatter: Callable | None = ..., na_rep: _str = ...
     ) -> list[_str]: ...
-    def to_flat_index(self): ...
     def to_series(
         self, index: Index | None = None, name: Hashable | None = None
     ) -> Series[S1]: ...
-    def to_frame(self, index: bool = True, name=...) -> DataFrame: ...
+    def to_frame(self, index: bool = True, name: Hashable = ...) -> DataFrame: ...
     @property
     def name(self) -> Hashable | None: ...
     @name.setter
@@ -392,11 +389,17 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
     def names(self) -> list[Hashable | None]: ...
     @names.setter
     def names(self, names: SequenceNotStr[Hashable | None]) -> None: ...
-    def set_names(self, names, *, level=..., inplace: bool = ...): ...
+    def set_names(
+        self,
+        names: Hashable | Sequence[Hashable],
+        *,
+        level: Level | Sequence[Level] | None = None,
+        inplace: bool = False,
+    ) -> Self: ...
     @overload
-    def rename(self, name, *, inplace: Literal[False] = False) -> Self: ...
+    def rename(self, name: Hashable, *, inplace: Literal[False] = False) -> Self: ...
     @overload
-    def rename(self, name, *, inplace: Literal[True]) -> None: ...
+    def rename(self, name: Hashable, *, inplace: Literal[True]) -> None: ...
     @property
     def nlevels(self) -> int: ...
     def get_level_values(self, level: int | _str) -> Index: ...

pandas-stubs/core/indexes/multi.pyi

@@ -2,9 +2,11 @@ from collections.abc import (
     Callable,
     Hashable,
     Iterable,
+    Mapping,
     Sequence,
 )
 from typing import (
+    Any,
     final,
     overload,
 )
@@ -109,13 +111,13 @@ class MultiIndex(Index):
     def dropna(self, how: AnyAll = "any") -> Self: ...
     def get_level_values(self, level: str | int) -> Index: ...
     def unique(self, level=...): ...
-    def to_frame(
+    def to_frame(  # type: ignore[override]  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         index: bool = True,
         name: list[HashableT] = ...,
         allow_duplicates: bool = False,
     ) -> pd.DataFrame: ...
-    def to_flat_index(self): ...
+    def to_flat_index(self) -> Index: ...
     def remove_unused_levels(self): ...
     @property
     def nlevels(self) -> int: ...
@@ -163,3 +165,10 @@ class MultiIndex(Index):
     def insert(self, loc, item): ...
     def delete(self, loc): ...
     def isin(self, values, level=...) -> np_1darray_bool: ...
+    def set_names(
+        self,
+        names: Hashable | Sequence[Hashable] | Mapping[Any, Hashable],
+        *,
+        level: Level | Sequence[Level] | None = None,
+        inplace: bool = False,
+    ) -> Self: ...

tests/indexes/test_indexes.py

@@ -1523,3 +1523,26 @@ def test_datetimeindex_where() -> None:
 
     val_range = pd.RangeIndex(2).where(pd.Series([True, False]), 3)
     check(assert_type(val_range, pd.Index), pd.RangeIndex)
+
+
+def test_index_set_names() -> None:
+    """Test Index.where with multiple types of other GH1419."""
+    idx = pd.Index([1, 2])
+    check(
+        assert_type(idx.set_names("chinchilla"), "pd.Index[int]"), pd.Index, np.integer
+    )
+    check(assert_type(idx.set_names((0,)), "pd.Index[int]"), pd.Index, np.integer)
+    check(
+        assert_type(idx.set_names(["chinchilla"]), "pd.Index[int]"),
+        pd.Index,
+        np.integer,
+    )
+
+    mi = pd.MultiIndex.from_arrays([[1, 2, 3], [4, 5, 6]], names=["elk", "owl"])
+    check(assert_type(mi.set_names(["beluga", "pig"]), pd.MultiIndex), pd.MultiIndex)
+    check(
+        assert_type(mi.set_names({"elk": "beluga", "owl": "pig"}), pd.MultiIndex),
+        pd.MultiIndex,
+    )
+    mi = cast("pd.MultiIndex", pd.Index([(1,)]))
+    check(assert_type(mi.set_names(1), pd.MultiIndex), pd.MultiIndex, tuple)