From 3aa3e8ff97ebae2150538519aa1c544a6daaf586 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Tue, 29 Dec 2020 12:29:47 -0800 Subject: [PATCH] TYP: ExtensionIndex (#38783) --- pandas/core/arrays/_mixins.py | 10 ++++++++-- pandas/core/indexes/extension.py | 9 ++++----- 2 files changed, 12 insertions(+), 7 deletions(-) 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)