type Index view and drop

Author: MarcoGorelli

pandas-stubs/_typing.pyi

@@ -950,6 +950,8 @@ np_1darray_dt: TypeAlias = np_1darray[np.datetime64]
 np_1darray_td: TypeAlias = np_1darray[np.timedelta64]
 np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[GenericT]]
 
+NDArrayT = TypeVar("NDArrayT", bound=np.ndarray)
+
 DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic])
 KeysArgType: TypeAlias = Any
 ListLikeT = TypeVar("ListLikeT", bound=ListLike)

pandas-stubs/core/indexes/base.pyi

@@ -102,6 +102,7 @@ from pandas._typing import (
     Level,
     MaskType,
     NaPosition,
+    NDArrayT,
     NumpyFloatNot16DtypeArg,
     PandasAstypeFloatDtypeArg,
     PandasFloatDtypeArg,
@@ -374,7 +375,12 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
     def dtype(self) -> DtypeObj: ...
     @final
     def ravel(self, order: _str = "C") -> Self: ...
-    def view(self, cls=...): ...
+    @overload
+    def view(self, cls: None = None) -> Self: ...
+    @overload
+    def view(self, cls: type[NDArrayT]) -> NDArrayT: ...
+    @overload
+    def view(self, cls: Dtype) -> ArrayLike: ...
     @overload
     def astype(
         self,
@@ -596,7 +602,11 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
     def insert(self, loc: int, item: S1) -> Self: ...
     @overload
     def insert(self, loc: int, item: object) -> Index: ...
-    def drop(self, labels, errors: IgnoreRaise = "raise") -> Self: ...
+    def drop(
+        self,
+        labels: IndexOpsMixin | np_1darray | Iterable[Hashable],
+        errors: IgnoreRaise = "raise",
+    ) -> Self: ...
     @property
     def shape(self) -> tuple[int, ...]: ...
     # Extra methods from old stubs

pyproject.toml

@@ -237,7 +237,11 @@ ignore = [
   # TODO: remove when _libs is fully typed
   "ANN001", "ANN201", "ANN204", "ANN206",
 ]
-"*base.pyi" = [
+"*core/base.pyi" = [
+  # TODO: remove when base.pyi's are fully typed
+  "ANN001", "ANN201", "ANN204", "ANN206",
+]
+"*excel/_base.pyi" = [
   # TODO: remove when base.pyi's are fully typed
   "ANN001", "ANN201", "ANN204", "ANN206",
 ]

tests/indexes/test_indexes.py

@@ -22,6 +22,7 @@
     assert_type,
 )
 
+from pandas._typing import ArrayLike  # noqa: F401
 from pandas._typing import Dtype  # noqa: F401
 from pandas._typing import Scalar  # noqa: F401
 
@@ -1667,3 +1668,28 @@ def test_index_slice_locs() -> None:
     start, end = idx.slice_locs(0, 1)
     check(assert_type(start, np.intp | int), np.integer)
     check(assert_type(end, np.intp | int), int)
+
+
+def test_index_view() -> None:
+    ind = pd.Index([1, 2])
+    check(assert_type(ind.view("int64"), ArrayLike), np_1darray_int64)
+    check(assert_type(ind.view(), "pd.Index[int]"), pd.Index)
+    check(assert_type(ind.view(np.ndarray), np.ndarray), np.ndarray)
+
+    class MyArray(np.ndarray): ...
+
+    check(assert_type(ind.view(MyArray), MyArray), MyArray)
+
+
+def test_index_drop() -> None:
+    ind = pd.Index([1, 2, 3])
+    check(assert_type(ind.drop([1, 2]), "pd.Index[int]"), pd.Index, np.integer)
+    check(
+        assert_type(ind.drop(pd.Index([1, 2])), "pd.Index[int]"), pd.Index, np.integer
+    )
+    check(
+        assert_type(ind.drop(pd.Series([1, 2])), "pd.Index[int]"), pd.Index, np.integer
+    )
+    check(
+        assert_type(ind.drop(np.array([1, 2])), "pd.Index[int]"), pd.Index, np.integer
+    )