From 3aae976e825d10fc66ada3457d0e526f04b10130 Mon Sep 17 00:00:00 2001 From: Brock Date: Tue, 29 Dec 2020 08:10:24 -0800 Subject: [PATCH 1/2] TYP: ExtensionIndex --- 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 b6938931e86af..3a17e8926359f 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): """ 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) From 0b62f090773f040670fba7e2e47b90d8e76fe50b Mon Sep 17 00:00:00 2001 From: Brock Date: Tue, 29 Dec 2020 09:50:22 -0800 Subject: [PATCH 2/2] TYP: putmask return type --- pandas/core/arrays/_mixins.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index 3a17e8926359f..4d4df96953196 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: NDArrayBackedExtensionArrayT, mask: np.ndarray, value): + def putmask(self: NDArrayBackedExtensionArrayT, mask: np.ndarray, value) -> None: """ Analogue to np.putmask(self, mask, value)