ENH: implement ExtensionArray.__array_function__ (GH#26380)

@@ -28,6 +28,7 @@ enhancement2  Other enhancements ^^^^^^^^^^^^^^^^^^+- :func:`numpy.delete` on an :class:`ExtensionArray`, and :func:`numpy.concatenate` on :class:`ExtensionArray` objects of the same dtype, now return an :class:`ExtensionArray` of that dtype instead of a NumPy array (:issue:`26380`) - :func:`read_csv` and :func:`read_table` with ``sep=None`` now sniff the separator with the C engine instead of falling back to the slower python engine (:issue:`9645`) - :func:`read_csv` with the ``"c"`` and ``"python"`` engines now accepts a :class:`DatetimeTZDtype` in ``dtype`` (:issue:`24542`) - :meth:`.DataFrameGroupBy.rank` and :meth:`.SeriesGroupBy.rank` now accept a ``numeric_only`` argument, matching :meth:`DataFrame.rank` (:issue:`44438`)@@ -88,6 +89,7 @@ See :ref:`install.dependencies` and :ref:`install.optional_dependencies` for mor Other API changes ^^^^^^^^^^^^^^^^^ - :class:`.Styler` now displays missing float values as ``NaN`` rather than ``nan`` by default, matching :class:`DataFrame` (:issue:`12710`)+- :class:`ExtensionArray` now implements ``__array_function__``. Code expecting :func:`numpy.concatenate` or :func:`numpy.delete` to return a NumPy array for :class:`ExtensionArray` input, including ``_concat_same_type`` and ``delete`` implementations in :class:`ExtensionArray` subclasses, may need updating, as may code mixing an :class:`ExtensionArray` with another library's ``__array_function__`` arrays (:issue:`26380`) - :class:`Index` now displays missing float values as ``NaN`` rather than ``nan`` in its repr, matching :class:`Series` and :class:`DataFrame` (:issue:`64733`)  .. ---------------------------------------------------------------------------

@@ -9,6 +9,7 @@  from __future__ import annotations +import contextvars import operator from typing import (     TYPE_CHECKING,@@ -2952,6 +2953,38 @@ def round(self, decimals: int = 0) -> Self:         ]         return type(self)._from_sequence(rounded, dtype=self.dtype) +    def __array_function__(self, func, types, args, kwargs):+        for typ in types:+            # defer to other types that override __array_function__; a subclass+            #  delegating here with super() is not "other"+            if not isinstance(self, typ) and typ.__array_function__ not in (+                ExtensionArray.__array_function__,+                np.ndarray.__array_function__,+            ):+                return NotImplemented++        override = None+        if not _IN_ARRAY_FUNCTION.get():+            override = _ARRAY_FUNCTION_OVERRIDES.get(func)+        implementation = getattr(func, "_implementation", None)+        if override is None and implementation is None:+            return NotImplemented++        # numpy calls made while handling this one (e.g. np.setxor1d calling+        #  np.concatenate) skip the overrides+        token = _IN_ARRAY_FUNCTION.set(True)+        try:+            if override is not None:+                result = override(*args, **kwargs)+                if result is not NotImplemented:+                    return result+            if implementation is None:+                return NotImplemented+            # numpy's own default, see ndarray.__array_function__ in NEP 18+            return implementation(*args, **kwargs)+        finally:+            _IN_ARRAY_FUNCTION.reset(token)+     def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):         # the fallback at the end of this method np.asarray()s the values: M8/m8         #  survive to be truth-tested, the rest flatten to object@@ -3235,6 +3268,53 @@ def _groupby_first_last(         return self.take(result_indices, allow_fill=True)  +# -----------------------------------------------------------------------------+# __array_function__ helpers++# True while ExtensionArray.__array_function__ is handling a call+_IN_ARRAY_FUNCTION = contextvars.ContextVar("_IN_ARRAY_FUNCTION", default=False)+++# Overrides so these keep the extension dtype. Returning NotImplemented falls+#  back to numpy's implementation.+++def _delete(arr, obj, axis=None):+    # the ExtensionArray may be obj rather than arr+    if (+        not isinstance(arr, ExtensionArray)+        or arr.ndim != 1+        or axis not in (None, 0, -1)+    ):+        return NotImplemented+    return arr.delete(obj)+++def _concatenate(arrays, axis=0, out=None, dtype=None, casting="same_kind"):+    if (+        not isinstance(arrays, (list, tuple))+        or not arrays+        or axis not in (None, 0, -1)+        or out is not None+        or dtype is not None+        or casting != "same_kind"+    ):+        return NotImplemented+    cls = type(arrays[0])+    if not issubclass(cls, ExtensionArray) or not all(+        type(arr) is cls and arr.ndim == 1 and arr.dtype == arrays[0].dtype+        for arr in arrays+    ):+        return NotImplemented+    return cls._concat_same_type(arrays)+++_ARRAY_FUNCTION_OVERRIDES: dict[Callable[..., Any], Callable[..., Any]] = {+    np.delete: _delete,+    np.concatenate: _concatenate,+}++ class ExtensionArrayNaResult(ExtensionArray):     @overload     def any(self, *, skipna: Literal[True] = ...) -> bool: ...

@@ -0,0 +1,160 @@+"""+Tests for ExtensionArray.__array_function__ beyond the extension base tests.+"""++import numpy as np+import pytest++import pandas as pd+import pandas._testing as tm+from pandas.arrays import (+    NumpyExtensionArray,+    SparseArray,+)+++def test_delete_with_extension_array_obj():+    # GH#26380 the ExtensionArray is the indexer, not the array+    result = np.delete(np.arange(3), pd.array([0]))+    tm.assert_numpy_array_equal(result, np.array([1, 2]))++    result = pd.Index([1, 2, 3]).delete(pd.array([0]))+    tm.assert_index_equal(result, pd.Index([2, 3]))+++@pytest.mark.parametrize("axis", [0, -1, None])+def test_axis_spellings(axis):+    arr = pd.array([1, None, 3], dtype="Int64")+    result = np.delete(arr, 0, axis=axis)+    tm.assert_extension_array_equal(result, arr[1:])++    result = np.concatenate([arr, arr[:1]], axis=axis)+    tm.assert_extension_array_equal(result, arr.take([0, 1, 2, 0]))+++@pytest.mark.parametrize(+    "kwargs",+    [{"dtype": object}, {"axis": None, "dtype": object}, {"casting": "unsafe"}],+)+def test_concatenate_with_kwargs_falls_back(kwargs):+    arr = pd.array([1, None], dtype="Int64")+    result = np.concatenate([arr, arr], **kwargs)+    expected = np.concatenate([np.asarray(arr)] * 2, **kwargs)+    tm.assert_numpy_array_equal(result, expected)+++def test_concatenate_out_falls_back():+    arr = pd.array([1, None], dtype="Int64")+    out = np.empty(4, dtype=object)+    result = np.concatenate([arr, arr], out=out)+    assert result is out+    expected = np.concatenate([np.asarray(arr)] * 2, out=np.empty(4, dtype=object))+    tm.assert_numpy_array_equal(out, expected)+++@pytest.mark.parametrize(+    "other",+    [+        pd.array([1.5], dtype="Float64"),+        pd.Categorical([1]),+        np.array([1]),+    ],+)+def test_concatenate_mismatched_falls_back(other):+    arr = pd.array([1, None], dtype="Int64")+    result = np.concatenate([arr, other])+    expected = np.concatenate([np.asarray(arr), np.asarray(other)])+    tm.assert_numpy_array_equal(result, expected)+++@pytest.mark.parametrize(+    "left, right",+    [+        (pd.Categorical(["a"]), pd.Categorical(["b"])),+        (+            pd.array(pd.date_range("2000", periods=2)),+            pd.array(pd.date_range("2000", periods=2, tz="UTC")),+        ),+    ],+)+def test_concatenate_same_type_different_dtype_falls_back(left, right):+    result = np.concatenate([left, right])+    expected = np.concatenate([np.asarray(left), np.asarray(right)])+    tm.assert_numpy_array_equal(result, expected)+++def test_2d_falls_back():+    arr = pd.date_range("2000", periods=4)._data.reshape(2, 2)+    expected = np.asarray(arr)++    result = np.delete(arr, 0, axis=0)+    tm.assert_numpy_array_equal(result, np.delete(expected, 0, axis=0))++    result = np.concatenate([arr, arr])+    tm.assert_numpy_array_equal(result, np.concatenate([expected, expected]))+++class ArrayFunctionOverride:+    def __array_function__(self, func, types, args, kwargs):+        return "override"+++def test_defers_to_other_array_function():+    arr = pd.array([1, 2], dtype="Int64")+    other = ArrayFunctionOverride()+    assert np.concatenate([arr, other]) == "override"+    assert np.concatenate([other, arr]) == "override"+++def test_defers_to_extension_array_subclass_override():+    class SubclassOverride(NumpyExtensionArray):+        def __array_function__(self, func, types, args, kwargs):+            return "override"++    arr = pd.array([1, 2], dtype="Int64")+    other = SubclassOverride(np.array([3, 4]))+    assert np.concatenate([arr, other]) == "override"+++def test_subclass_override_delegating_to_super():+    class DelegatingOverride(NumpyExtensionArray):+        def __array_function__(self, func, types, args, kwargs):+            return super().__array_function__(func, types, args, kwargs)++    arr = DelegatingOverride(np.array([1, 2, 3]))+    assert np.mean(arr) == 2.0+    result = np.delete(arr, 0)+    tm.assert_extension_array_equal(result, DelegatingOverride(np.array([2, 3])))+++@pytest.mark.parametrize(+    "left, right",+    [+        (pd.array([3, None, 2], dtype="Int64"), pd.array([2, None], dtype="Int64")),+        (SparseArray([3, 1, 2]), SparseArray([2, 5])),+    ],+)+def test_nested_numpy_calls_unchanged(left, right):+    # GH#26380 np.setxor1d calls np.concatenate internally, which must keep+    #  returning an ndarray there+    result = np.setxor1d(left, right, assume_unique=True)+    expected = np.setxor1d(np.asarray(left), np.asarray(right), assume_unique=True)+    tm.assert_numpy_array_equal(result, expected)+++def test_concat_same_type_calling_np_concatenate():+    # GH#26380 a subclass implementing _concat_same_type with np.concatenate+    #  must not recurse endlessly, whether called via numpy or directly+    class ConcatWithNumpy(NumpyExtensionArray):+        @classmethod+        def _concat_same_type(cls, to_concat):+            return cls(np.concatenate(to_concat))++    arr = ConcatWithNumpy(np.arange(3))+    expected = ConcatWithNumpy(np.array([0, 1, 2, 0, 1, 2]))++    result = np.concatenate([arr, arr])+    tm.assert_extension_array_equal(result, expected)++    result = ConcatWithNumpy._concat_same_type([arr, arr])+    tm.assert_extension_array_equal(result, expected)

@@ -815,12 +815,13 @@ def test_repeat_raises(self, data, repeats, kwargs, error, msg, use_numpy):             else:                 data.repeat(repeats, **kwargs) -    def test_delete(self, data):-        result = data.delete(0)+    def test_delete(self, data, use_numpy):+        # GH#26380 np.delete keeps the dtype+        result = np.delete(data, 0) if use_numpy else data.delete(0)         expected = data[1:]         tm.assert_extension_array_equal(result, expected) -        result = data.delete([1, 3])+        result = np.delete(data, [1, 3]) if use_numpy else data.delete([1, 3])         expected = data._concat_same_type([data[[0]], data[[2]], data[4:]])         tm.assert_extension_array_equal(result, expected) 

@@ -33,6 +33,12 @@ def test_concat(self, data, in_frame):             assert isinstance(result._mgr.blocks[0], EABackedBlock)         assert isinstance(result._mgr.blocks[0].values, ExtensionArray) +    def test_np_concatenate(self, data):+        # GH#26380 np.concatenate on matching dtypes keeps the dtype+        result = np.concatenate([data, data[:3]])+        expected = data._concat_same_type([data, data[:3]])+        tm.assert_extension_array_equal(result, expected)+     @pytest.mark.parametrize("in_frame", [True, False])     def test_concat_all_na_block(self, data_missing, in_frame):         valid_block = pd.Series(data_missing.take([1, 1]), index=range(2))