ENH: NDArrayBackedExtensionArray.__array_function__

Author: jbrockmendel

pandas/core/arrays/_mixins.py

@@ -304,6 +304,48 @@ def __repr__(self) -> str:
     # ------------------------------------------------------------------------
     # __array_function__ methods
 
+    def __array_function__(self, func, types, args, kwargs):
+
+        if not args:
+            # I dont think this is possible is it?
+            raise NotImplementedError
+
+        if args[0] is self:
+
+            if func in [np.delete, np.repeat, np.atleast_2d]:
+                res_data = func(self._ndarray, *args[1:], **kwargs)
+                return self._from_backing_data(res_data)
+
+            # TODO: do we need to convert args to kwargs to ensure nv checks
+            #  are correct?
+            if func is np.amin:
+                return self.min(*args[1:], **kwargs)
+            if func is np.amax:
+                return self.max(*args[1:], **kwargs)
+
+            if func is np.sum:
+                # Need to do explicitly otherise np.sum(TimedeltaArray)
+                #  doesnt wrap in Timedelta.
+                return self.sum(*args[1:], **kwargs)
+
+            if func is np.argsort:
+                if len(args) > 1:
+                    # try to make sure that we are passing kwargs along correclty
+                    raise NotImplementedError
+                return self.argsort(*args[1:], **kwargs)
+
+        if not any(x is self for x in args):
+            # e.g. np.conatenate we get args[0] is a tuple containing self
+            largs = list(args)
+            for i, arg in enumerate(largs):
+                if isinstance(arg, (list, tuple)):
+                    arg = type(arg)(x if x is not self else np.asarray(x) for x in arg)
+                    largs[i] = arg
+            args = tuple(largs)
+
+        args = [x if x is not self else np.asarray(x) for x in args]
+        return func(*args, **kwargs)
+
     def putmask(self, mask, value):
         """
         Analogue to np.putmask(self, mask, value)

pandas/tests/arrays/test_ndarray_backed.py

@@ -0,0 +1,105 @@
+"""
+Tests for EA subclasses subclassing NDArrayBackedExtensionArray
+"""
+
+import numpy as np
+import pytest
+
+from pandas import date_range
+import pandas._testing as tm
+from pandas.core.arrays import Categorical, PandasArray
+
+
+class ArrayFunctionTests:
+    # Tests for subclasses that do not explicitly support 2D yet.
+    def test_delete_no_axis(self, array):
+        # with no axis, operates on flattened version
+        result = np.delete(array, 1)
+
+        backing = np.delete(array._ndarray.ravel(), 1)
+        expected = array._from_backing_data(backing)
+        tm.assert_equal(result, expected)
+
+    def test_repeat(self, array):
+        result = np.repeat(array, 2)
+
+        backing = np.repeat(array._ndarray.ravel(), 2)
+        expected = array._from_backing_data(backing)
+        tm.assert_equal(result, expected)
+
+
+class ArrayFunctionTests2D(ArrayFunctionTests):
+    @pytest.mark.parametrize("axis", [0, 1])
+    def test_delete_axis(self, array, axis):
+        result = np.delete(array, 1, axis=axis)
+        if axis == 0:
+            assert result.shape == (array.shape[0] - 1, array.shape[1])
+        else:
+            assert result.shape == (array.shape[0], array.shape[1] - 1)
+
+        backing = np.delete(array._ndarray, 1, axis=axis)
+        expected = array._from_backing_data(backing)
+        tm.assert_equal(result, expected)
+
+        # axis as an arg instead of as a kwarg
+        result = np.delete(array, 1, axis)
+        tm.assert_equal(result, expected)
+
+    @pytest.mark.parametrize("axis", [0, 1])
+    def test_repeat_axis(self, array, axis):
+        result = np.repeat(array, 2, axis=axis)
+
+        backing = np.repeat(array._ndarray, 2, axis=axis)
+        expected = array._from_backing_data(backing)
+        tm.assert_equal(result, expected)
+
+        # axis as an arg instead of a kwarg
+        result = np.repeat(array, 2, axis)
+        tm.assert_equal(result, expected)
+
+    def test_atleast_2d(self, array):
+        result = np.atleast_2d(array)
+
+        assert result.ndim >= 2
+
+        if array.ndim == 1:
+            assert result.shape == (1, array.size)
+        else:
+            assert result.shape == array.shape
+
+
+class TestDatetimeArray(ArrayFunctionTests2D):
+    @pytest.fixture(params=[1, 2])
+    def array(self):
+        dti = date_range("1994-05-12", periods=12, tz="US/Pacific")
+        dta = dti._data.reshape(3, 4)
+        return dta
+
+
+class TestTimedeltaArray(ArrayFunctionTests2D):
+    @pytest.fixture
+    def array(self):
+        dti = date_range("1994-05-12", periods=12, tz="US/Pacific")
+        dta = dti._data.reshape(3, 4)
+        return dta - dta[0, 0]
+
+
+class TestPeriodArray(ArrayFunctionTests2D):
+    @pytest.fixture
+    def array(self):
+        dti = date_range("1994-05-12", periods=12)
+        pa = dti._data.to_period("D")
+        return pa.reshape(3, 4)
+
+
+class TestPandasArray(ArrayFunctionTests):
+    @pytest.fixture
+    def array(self):
+        return PandasArray(np.arange(12))
+
+
+class TestCategorical(ArrayFunctionTests):
+    @pytest.fixture
+    def array(self):
+        dti = date_range("1994-05-12", periods=12, tz="US/Pacific")
+        return Categorical(dti)