TST: port Dim2CompatTests (#39880)

Author: jbrockmendel

pandas/core/ops/mask_ops.py

@@ -179,6 +179,6 @@ def kleene_and(
     return result, mask
 
 
-def raise_for_nan(value, method):
+def raise_for_nan(value, method: str):
     if lib.is_float(value) and np.isnan(value):
         raise ValueError(f"Cannot perform logical '{method}' with floating NaN")

pandas/tests/extension/base/__init__.py

@@ -43,6 +43,7 @@ class TestMyDtype(BaseDtypeTests):
 """
 from pandas.tests.extension.base.casting import BaseCastingTests  # noqa
 from pandas.tests.extension.base.constructors import BaseConstructorsTests  # noqa
+from pandas.tests.extension.base.dim2 import Dim2CompatTests  # noqa
 from pandas.tests.extension.base.dtype import BaseDtypeTests  # noqa
 from pandas.tests.extension.base.getitem import BaseGetitemTests  # noqa
 from pandas.tests.extension.base.groupby import BaseGroupbyTests  # noqa

pandas/tests/extension/base/dim2.py

@@ -0,0 +1,217 @@
+"""
+Tests for 2D compatibility.
+"""
+import numpy as np
+import pytest
+
+from pandas.compat import np_version_under1p17
+
+import pandas as pd
+from pandas.core.arrays import (
+    FloatingArray,
+    IntegerArray,
+)
+from pandas.tests.extension.base.base import BaseExtensionTests
+
+
+def maybe_xfail_masked_reductions(arr, request):
+    if (
+        isinstance(arr, (FloatingArray, IntegerArray))
+        and np_version_under1p17
+        and arr.ndim == 2
+    ):
+        mark = pytest.mark.xfail(reason="masked_reductions does not implement")
+        request.node.add_marker(mark)
+
+
+class Dim2CompatTests(BaseExtensionTests):
+    def test_take_2d(self, data):
+        arr2d = data.reshape(-1, 1)
+
+        result = arr2d.take([0, 0, -1], axis=0)
+
+        expected = data.take([0, 0, -1]).reshape(-1, 1)
+        self.assert_extension_array_equal(result, expected)
+
+    def test_repr_2d(self, data):
+        # this could fail in a corner case where an element contained the name
+        res = repr(data.reshape(1, -1))
+        assert res.count(f"<{type(data).__name__}") == 1
+
+        res = repr(data.reshape(-1, 1))
+        assert res.count(f"<{type(data).__name__}") == 1
+
+    def test_reshape(self, data):
+        arr2d = data.reshape(-1, 1)
+        assert arr2d.shape == (data.size, 1)
+        assert len(arr2d) == len(data)
+
+        arr2d = data.reshape((-1, 1))
+        assert arr2d.shape == (data.size, 1)
+        assert len(arr2d) == len(data)
+
+        with pytest.raises(ValueError):
+            data.reshape((data.size, 2))
+        with pytest.raises(ValueError):
+            data.reshape(data.size, 2)
+
+    def test_getitem_2d(self, data):
+        arr2d = data.reshape(1, -1)
+
+        result = arr2d[0]
+        self.assert_extension_array_equal(result, data)
+
+        with pytest.raises(IndexError):
+            arr2d[1]
+
+        with pytest.raises(IndexError):
+            arr2d[-2]
+
+        result = arr2d[:]
+        self.assert_extension_array_equal(result, arr2d)
+
+        result = arr2d[:, :]
+        self.assert_extension_array_equal(result, arr2d)
+
+        result = arr2d[:, 0]
+        expected = data[[0]]
+        self.assert_extension_array_equal(result, expected)
+
+        # dimension-expanding getitem on 1D
+        result = data[:, np.newaxis]
+        self.assert_extension_array_equal(result, arr2d.T)
+
+    def test_iter_2d(self, data):
+        arr2d = data.reshape(1, -1)
+
+        objs = list(iter(arr2d))
+        assert len(objs) == arr2d.shape[0]
+
+        for obj in objs:
+            assert isinstance(obj, type(data))
+            assert obj.dtype == data.dtype
+            assert obj.ndim == 1
+            assert len(obj) == arr2d.shape[1]
+
+    def test_concat_2d(self, data):
+        left = data.reshape(-1, 1)
+        right = left.copy()
+
+        # axis=0
+        result = left._concat_same_type([left, right], axis=0)
+        expected = data._concat_same_type([data, data]).reshape(-1, 1)
+        self.assert_extension_array_equal(result, expected)
+
+        # axis=1
+        result = left._concat_same_type([left, right], axis=1)
+        expected = data.repeat(2).reshape(-1, 2)
+        self.assert_extension_array_equal(result, expected)
+
+        # axis > 1 -> invalid
+        with pytest.raises(ValueError):
+            left._concat_same_type([left, right], axis=2)
+
+    @pytest.mark.parametrize("method", ["mean", "median", "var", "std", "sum", "prod"])
+    def test_reductions_2d_axis_none(self, data, method, request):
+        if not hasattr(data, method):
+            pytest.skip("test is not applicable for this type/dtype")
+
+        arr2d = data.reshape(1, -1)
+        maybe_xfail_masked_reductions(arr2d, request)
+
+        err_expected = None
+        err_result = None
+        try:
+            expected = getattr(data, method)()
+        except Exception as err:
+            # if the 1D reduction is invalid, the 2D reduction should be as well
+            err_expected = err
+            try:
+                result = getattr(arr2d, method)(axis=None)
+            except Exception as err2:
+                err_result = err2
+
+        else:
+            result = getattr(arr2d, method)(axis=None)
+
+        if err_result is not None or err_expected is not None:
+            assert type(err_result) == type(err_expected)
+            return
+
+        assert result == expected  # TODO: or matching NA
+
+    @pytest.mark.parametrize("method", ["mean", "median", "var", "std", "sum", "prod"])
+    def test_reductions_2d_axis0(self, data, method, request):
+        if not hasattr(data, method):
+            pytest.skip("test is not applicable for this type/dtype")
+
+        arr2d = data.reshape(1, -1)
+        maybe_xfail_masked_reductions(arr2d, request)
+
+        kwargs = {}
+        if method == "std":
+            # pass ddof=0 so we get all-zero std instead of all-NA std
+            kwargs["ddof"] = 0
+
+        try:
+            result = getattr(arr2d, method)(axis=0, **kwargs)
+        except Exception as err:
+            try:
+                getattr(data, method)()
+            except Exception as err2:
+                assert type(err) == type(err2)
+                return
+            else:
+                raise AssertionError("Both reductions should raise or neither")
+
+        if method in ["mean", "median", "sum", "prod"]:
+            # std and var are not dtype-preserving
+            expected = data
+            if method in ["sum", "prod"] and data.dtype.kind in ["i", "u"]:
+                # FIXME: kludge
+                if data.dtype.kind == "i":
+                    dtype = pd.Int64Dtype
+                else:
+                    dtype = pd.UInt64Dtype
+
+                expected = data.astype(dtype)
+                if type(expected) != type(data):
+                    mark = pytest.mark.xfail(
+                        reason="IntegerArray.astype is broken GH#38983"
+                    )
+                    request.node.add_marker(mark)
+                assert type(expected) == type(data), type(expected)
+                assert dtype == expected.dtype
+
+            self.assert_extension_array_equal(result, expected)
+        elif method == "std":
+            self.assert_extension_array_equal(result, data - data)
+        # punt on method == "var"
+
+    @pytest.mark.parametrize("method", ["mean", "median", "var", "std", "sum", "prod"])
+    def test_reductions_2d_axis1(self, data, method, request):
+        if not hasattr(data, method):
+            pytest.skip("test is not applicable for this type/dtype")
+
+        arr2d = data.reshape(1, -1)
+        maybe_xfail_masked_reductions(arr2d, request)
+
+        try:
+            result = getattr(arr2d, method)(axis=1)
+        except Exception as err:
+            try:
+                getattr(data, method)()
+            except Exception as err2:
+                assert type(err) == type(err2)
+                return
+            else:
+                raise AssertionError("Both reductions should raise or neither")
+
+        # not necesarrily type/dtype-preserving, so weaker assertions
+        assert result.shape == (1,)
+        expected_scalar = getattr(data, method)()
+        if pd.isna(result[0]):
+            # TODO: require matching NA
+            assert pd.isna(expected_scalar), expected_scalar
+        else:
+            assert result[0] == expected_scalar

pandas/tests/extension/test_datetime.py

@@ -235,3 +235,7 @@ class TestGroupby(BaseDatetimeTests, base.BaseGroupbyTests):
 
 class TestPrinting(BaseDatetimeTests, base.BasePrintingTests):
     pass
+
+
+class Test2DCompat(BaseDatetimeTests, base.Dim2CompatTests):
+    pass

pandas/tests/extension/test_numpy.py

@@ -415,3 +415,7 @@ def test_setitem_loc_iloc_slice(self, data):
 @skip_nested
 class TestParsing(BaseNumPyTests, base.BaseParsingTests):
     pass
+
+
+class Test2DCompat(BaseNumPyTests, base.Dim2CompatTests):
+    pass

pandas/tests/extension/test_period.py

@@ -184,3 +184,7 @@ class TestParsing(BasePeriodTests, base.BaseParsingTests):
     @pytest.mark.parametrize("engine", ["c", "python"])
     def test_EA_types(self, engine, data):
         super().test_EA_types(engine, data)
+
+
+class Test2DCompat(BasePeriodTests, base.Dim2CompatTests):
+    pass