diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index 57df9ed9ed2d7..ad0bf76b0556b 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -369,8 +369,16 @@ def where( 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) + def delete( + self: NDArrayBackedExtensionArrayT, loc, axis: int = 0 + ) -> NDArrayBackedExtensionArrayT: + res_values = np.delete(self._ndarray, loc, axis=axis) + return self._from_backing_data(res_values) + + def swapaxes( + self: NDArrayBackedExtensionArrayT, axis1, axis2 + ) -> NDArrayBackedExtensionArrayT: + res_values = self._ndarray.swapaxes(axis1, axis2) return self._from_backing_data(res_values) # ------------------------------------------------------------------------ diff --git a/pandas/tests/extension/base/dim2.py b/pandas/tests/extension/base/dim2.py index c6455ce15533a..fbe2537e8a7bf 100644 --- a/pandas/tests/extension/base/dim2.py +++ b/pandas/tests/extension/base/dim2.py @@ -25,6 +25,26 @@ def maybe_xfail_masked_reductions(arr, request): class Dim2CompatTests(BaseExtensionTests): + def test_swapaxes(self, data): + arr2d = data.repeat(2).reshape(-1, 2) + + result = arr2d.swapaxes(0, 1) + expected = arr2d.T + self.assert_extension_array_equal(result, expected) + + def test_delete_2d(self, data): + arr2d = data.repeat(3).reshape(-1, 3) + + # axis = 0 + result = arr2d.delete(1, axis=0) + expected = data.delete(1).repeat(3).reshape(-1, 3) + self.assert_extension_array_equal(result, expected) + + # axis = 1 + result = arr2d.delete(1, axis=1) + expected = data.repeat(2).reshape(-1, 2) + self.assert_extension_array_equal(result, expected) + def test_take_2d(self, data): arr2d = data.reshape(-1, 1)