diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index 2f292c8db025e..62c594d73a227 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -325,7 +325,7 @@ def __repr__(self) -> str: # ------------------------------------------------------------------------ # __array_function__ methods - def putmask(self, mask, value): + def putmask(self: NDArrayBackedExtensionArrayT, mask: np.ndarray, value) -> None: """ Analogue to np.putmask(self, mask, value) @@ -343,7 +343,9 @@ def putmask(self, mask, value): np.putmask(self._ndarray, mask, value) - def where(self, mask, value): + def where( + self: NDArrayBackedExtensionArrayT, mask: np.ndarray, value + ) -> NDArrayBackedExtensionArrayT: """ Analogue to np.where(mask, self, value) @@ -361,3 +363,7 @@ def where(self, mask, value): res_values = np.where(mask, self._ndarray, value) return self._from_backing_data(res_values) + + def delete(self: NDArrayBackedExtensionArrayT, loc) -> NDArrayBackedExtensionArrayT: + res_values = np.delete(self._ndarray, loc) + return self._from_backing_data(res_values) diff --git a/pandas/core/indexes/extension.py b/pandas/core/indexes/extension.py index 661adde44089c..f597175bf2ae2 100644 --- a/pandas/core/indexes/extension.py +++ b/pandas/core/indexes/extension.py @@ -333,7 +333,7 @@ class NDArrayBackedExtensionIndex(ExtensionIndex): def _get_engine_target(self) -> np.ndarray: return self._data._ndarray - def delete(self, loc): + def delete(self: _T, loc) -> _T: """ Make new Index with passed location(-s) deleted @@ -341,11 +341,10 @@ def delete(self, loc): ------- new_index : Index """ - new_vals = np.delete(self._data._ndarray, loc) - arr = self._data._from_backing_data(new_vals) + arr = self._data.delete(loc) return type(self)._simple_new(arr, name=self.name) - def insert(self, loc: int, item): + def insert(self: _T, loc: int, item) -> _T: """ Make new Index inserting new item at location. Follows Python list.append semantics for negative values. @@ -371,7 +370,7 @@ def insert(self, loc: int, item): return type(self)._simple_new(new_arr, name=self.name) @doc(Index.where) - def where(self, cond, other=None): + def where(self: _T, cond: np.ndarray, other=None) -> _T: res_values = self._data.where(cond, other) return type(self)._simple_new(res_values, name=self.name)