REF: de-duplicate MaskedArray tests (#44004)

* REF: share tests

* REF: de-duplicate _compare_other

Author: jbrockmendel

pandas/tests/arrays/boolean/test_comparison.py

@@ -4,7 +4,7 @@
 import pandas as pd
 import pandas._testing as tm
 from pandas.arrays import BooleanArray
-from pandas.tests.extension.base import BaseOpsUtil
+from pandas.tests.arrays.masked_shared import ComparisonOps
 
 
 @pytest.fixture
@@ -15,30 +15,12 @@ def data():
     )
 
 
-class TestComparisonOps(BaseOpsUtil):
-    def _compare_other(self, data, op_name, other):
-        op = self.get_op_from_name(op_name)
-
-        # array
-        result = pd.Series(op(data, other))
-        expected = pd.Series(op(data._data, other), dtype="boolean")
-        # propagate NAs
-        expected[data._mask] = pd.NA
-
-        tm.assert_series_equal(result, expected)
-
-        # series
-        s = pd.Series(data)
-        result = op(s, other)
-
-        expected = pd.Series(data._data)
-        expected = op(expected, other)
-        expected = expected.astype("boolean")
-        # propagate NAs
-        expected[data._mask] = pd.NA
+@pytest.fixture
+def dtype():
+    return pd.BooleanDtype()
 
-        tm.assert_series_equal(result, expected)
 
+class TestComparisonOps(ComparisonOps):
     def test_compare_scalar(self, data, all_compare_operators):
         op_name = all_compare_operators
         self._compare_other(data, op_name, True)
@@ -53,24 +35,8 @@ def test_compare_array(self, data, all_compare_operators):
         self._compare_other(data, op_name, other)
 
     @pytest.mark.parametrize("other", [True, False, pd.NA])
-    def test_scalar(self, other, all_compare_operators):
-        op = self.get_op_from_name(all_compare_operators)
-        a = pd.array([True, False, None], dtype="boolean")
-
-        result = op(a, other)
-
-        if other is pd.NA:
-            expected = pd.array([None, None, None], dtype="boolean")
-        else:
-            values = op(a._data, other)
-            expected = BooleanArray(values, a._mask, copy=True)
-        tm.assert_extension_array_equal(result, expected)
-
-        # ensure we haven't mutated anything inplace
-        result[0] = None
-        tm.assert_extension_array_equal(
-            a, pd.array([True, False, None], dtype="boolean")
-        )
+    def test_scalar(self, other, all_compare_operators, dtype):
+        ComparisonOps.test_scalar(self, other, all_compare_operators, dtype)
 
     def test_array(self, all_compare_operators):
         op = self.get_op_from_name(all_compare_operators)

pandas/tests/arrays/floating/test_comparison.py

@@ -1,86 +1,17 @@
-import numpy as np
 import pytest
 
 import pandas as pd
 import pandas._testing as tm
-from pandas.tests.extension.base import BaseOpsUtil
+from pandas.tests.arrays.masked_shared import (
+    ComparisonOps,
+    NumericOps,
+)
 
 
-class TestComparisonOps(BaseOpsUtil):
-    def _compare_other(self, data, op_name, other):
-        op = self.get_op_from_name(op_name)
-
-        # array
-        result = pd.Series(op(data, other))
-        expected = pd.Series(op(data._data, other), dtype="boolean")
-
-        # fill the nan locations
-        expected[data._mask] = pd.NA
-
-        tm.assert_series_equal(result, expected)
-
-        # series
-        s = pd.Series(data)
-        result = op(s, other)
-
-        expected = op(pd.Series(data._data), other)
-
-        # fill the nan locations
-        expected[data._mask] = pd.NA
-        expected = expected.astype("boolean")
-
-        tm.assert_series_equal(result, expected)
-
+class TestComparisonOps(NumericOps, ComparisonOps):
     @pytest.mark.parametrize("other", [True, False, pd.NA, -1.0, 0.0, 1])
-    def test_scalar(self, other, all_compare_operators):
-        op = self.get_op_from_name(all_compare_operators)
-        a = pd.array([1.0, 0.0, None], dtype="Float64")
-
-        result = op(a, other)
-
-        if other is pd.NA:
-            expected = pd.array([None, None, None], dtype="boolean")
-        else:
-            values = op(a._data, other)
-            expected = pd.arrays.BooleanArray(values, a._mask, copy=True)
-        tm.assert_extension_array_equal(result, expected)
-
-        # ensure we haven't mutated anything inplace
-        result[0] = pd.NA
-        tm.assert_extension_array_equal(a, pd.array([1.0, 0.0, None], dtype="Float64"))
-
-    def test_array(self, all_compare_operators):
-        op = self.get_op_from_name(all_compare_operators)
-        a = pd.array([0, 1, 2, None, None, None], dtype="Float64")
-        b = pd.array([0, 1, None, 0, 1, None], dtype="Float64")
-
-        result = op(a, b)
-        values = op(a._data, b._data)
-        mask = a._mask | b._mask
-
-        expected = pd.arrays.BooleanArray(values, mask)
-        tm.assert_extension_array_equal(result, expected)
-
-        # ensure we haven't mutated anything inplace
-        result[0] = pd.NA
-        tm.assert_extension_array_equal(
-            a, pd.array([0, 1, 2, None, None, None], dtype="Float64")
-        )
-        tm.assert_extension_array_equal(
-            b, pd.array([0, 1, None, 0, 1, None], dtype="Float64")
-        )
-
-    def test_compare_with_booleanarray(self, all_compare_operators):
-        op = self.get_op_from_name(all_compare_operators)
-        a = pd.array([True, False, None] * 3, dtype="boolean")
-        b = pd.array([0] * 3 + [1] * 3 + [None] * 3, dtype="Float64")
-        other = pd.array([False] * 3 + [True] * 3 + [None] * 3, dtype="boolean")
-        expected = op(a, other)
-        result = op(a, b)
-        tm.assert_extension_array_equal(result, expected)
-        expected = op(other, a)
-        result = op(b, a)
-        tm.assert_extension_array_equal(result, expected)
+    def test_scalar(self, other, all_compare_operators, dtype):
+        ComparisonOps.test_scalar(self, other, all_compare_operators, dtype)
 
     def test_compare_with_integerarray(self, all_compare_operators):
         op = self.get_op_from_name(all_compare_operators)
@@ -94,18 +25,6 @@ def test_compare_with_integerarray(self, all_compare_operators):
         result = op(b, a)
         tm.assert_extension_array_equal(result, expected)
 
-    def test_no_shared_mask(self, data):
-        result = data + 1
-        assert np.shares_memory(result._mask, data._mask) is False
-
-    def test_compare_to_string(self, dtype):
-        # GH 28930
-        s = pd.Series([1, None], dtype=dtype)
-        result = s == "a"
-        expected = pd.Series([False, pd.NA], dtype="boolean")
-
-        self.assert_series_equal(result, expected)
-
 
 def test_equals():
     # GH-30652

pandas/tests/arrays/integer/test_comparison.py

@@ -1,99 +1,20 @@
-import numpy as np
 import pytest
 
 import pandas as pd
-import pandas._testing as tm
-from pandas.tests.extension.base import BaseOpsUtil
+from pandas.tests.arrays.masked_shared import (
+    ComparisonOps,
+    NumericOps,
+)
 
 
-class TestComparisonOps(BaseOpsUtil):
-    def _compare_other(self, data, op_name, other):
-        op = self.get_op_from_name(op_name)
-
-        # array
-        result = pd.Series(op(data, other))
-        expected = pd.Series(op(data._data, other), dtype="boolean")
-
-        # fill the nan locations
-        expected[data._mask] = pd.NA
-
-        tm.assert_series_equal(result, expected)
-
-        # series
-        s = pd.Series(data)
-        result = op(s, other)
-
-        expected = op(pd.Series(data._data), other)
-
-        # fill the nan locations
-        expected[data._mask] = pd.NA
-        expected = expected.astype("boolean")
-
-        tm.assert_series_equal(result, expected)
-
+class TestComparisonOps(NumericOps, ComparisonOps):
     @pytest.mark.parametrize("other", [True, False, pd.NA, -1, 0, 1])
-    def test_scalar(self, other, all_compare_operators):
-        op = self.get_op_from_name(all_compare_operators)
-        a = pd.array([1, 0, None], dtype="Int64")
-
-        result = op(a, other)
-
-        if other is pd.NA:
-            expected = pd.array([None, None, None], dtype="boolean")
-        else:
-            values = op(a._data, other)
-            expected = pd.arrays.BooleanArray(values, a._mask, copy=True)
-        tm.assert_extension_array_equal(result, expected)
-
-        # ensure we haven't mutated anything inplace
-        result[0] = pd.NA
-        tm.assert_extension_array_equal(a, pd.array([1, 0, None], dtype="Int64"))
-
-    def test_array(self, all_compare_operators):
-        op = self.get_op_from_name(all_compare_operators)
-        a = pd.array([0, 1, 2, None, None, None], dtype="Int64")
-        b = pd.array([0, 1, None, 0, 1, None], dtype="Int64")
-
-        result = op(a, b)
-        values = op(a._data, b._data)
-        mask = a._mask | b._mask
-
-        expected = pd.arrays.BooleanArray(values, mask)
-        tm.assert_extension_array_equal(result, expected)
-
-        # ensure we haven't mutated anything inplace
-        result[0] = pd.NA
-        tm.assert_extension_array_equal(
-            a, pd.array([0, 1, 2, None, None, None], dtype="Int64")
-        )
-        tm.assert_extension_array_equal(
-            b, pd.array([0, 1, None, 0, 1, None], dtype="Int64")
-        )
-
-    def test_compare_with_booleanarray(self, all_compare_operators):
-        op = self.get_op_from_name(all_compare_operators)
-        a = pd.array([True, False, None] * 3, dtype="boolean")
-        b = pd.array([0] * 3 + [1] * 3 + [None] * 3, dtype="Int64")
-        other = pd.array([False] * 3 + [True] * 3 + [None] * 3, dtype="boolean")
-        expected = op(a, other)
-        result = op(a, b)
-        tm.assert_extension_array_equal(result, expected)
-
-    def test_no_shared_mask(self, data):
-        result = data + 1
-        assert np.shares_memory(result._mask, data._mask) is False
-
-    def test_compare_to_string(self, any_int_ea_dtype):
-        # GH 28930
-        s = pd.Series([1, None], dtype=any_int_ea_dtype)
-        result = s == "a"
-        expected = pd.Series([False, pd.NA], dtype="boolean")
-
-        self.assert_series_equal(result, expected)
+    def test_scalar(self, other, all_compare_operators, dtype):
+        ComparisonOps.test_scalar(self, other, all_compare_operators, dtype)
 
-    def test_compare_to_int(self, any_int_ea_dtype, all_compare_operators):
+    def test_compare_to_int(self, dtype, all_compare_operators):
         # GH 28930
-        s1 = pd.Series([1, None, 3], dtype=any_int_ea_dtype)
+        s1 = pd.Series([1, None, 3], dtype=dtype)
         s2 = pd.Series([1, None, 3], dtype="float")
 
         method = getattr(s1, all_compare_operators)

pandas/tests/arrays/masked_shared.py

@@ -0,0 +1,106 @@
+"""
+Tests shared by MaskedArray subclasses.
+"""
+import numpy as np
+
+import pandas as pd
+import pandas._testing as tm
+from pandas.tests.extension.base import BaseOpsUtil
+
+
+class ComparisonOps(BaseOpsUtil):
+    def _compare_other(self, data, op_name, other):
+        op = self.get_op_from_name(op_name)
+
+        # array
+        result = pd.Series(op(data, other))
+        expected = pd.Series(op(data._data, other), dtype="boolean")
+
+        # fill the nan locations
+        expected[data._mask] = pd.NA
+
+        tm.assert_series_equal(result, expected)
+
+        # series
+        ser = pd.Series(data)
+        result = op(ser, other)
+
+        expected = op(pd.Series(data._data), other)
+
+        # fill the nan locations
+        expected[data._mask] = pd.NA
+        expected = expected.astype("boolean")
+
+        tm.assert_series_equal(result, expected)
+
+    # subclass will override to parametrize 'other'
+    def test_scalar(self, other, all_compare_operators, dtype):
+        op = self.get_op_from_name(all_compare_operators)
+        left = pd.array([1, 0, None], dtype=dtype)
+
+        result = op(left, other)
+
+        if other is pd.NA:
+            expected = pd.array([None, None, None], dtype="boolean")
+        else:
+            values = op(left._data, other)
+            expected = pd.arrays.BooleanArray(values, left._mask, copy=True)
+        tm.assert_extension_array_equal(result, expected)
+
+        # ensure we haven't mutated anything inplace
+        result[0] = pd.NA
+        tm.assert_extension_array_equal(left, pd.array([1, 0, None], dtype=dtype))
+
+
+class NumericOps:
+    # Shared by IntegerArray and FloatingArray, not BooleanArray
+
+    def test_no_shared_mask(self, data):
+        result = data + 1
+        assert np.shares_memory(result._mask, data._mask) is False
+
+    def test_array(self, all_compare_operators, dtype):
+        op = self.get_op_from_name(all_compare_operators)
+
+        left = pd.array([0, 1, 2, None, None, None], dtype=dtype)
+        right = pd.array([0, 1, None, 0, 1, None], dtype=dtype)
+
+        result = op(left, right)
+        values = op(left._data, right._data)
+        mask = left._mask | right._mask
+
+        expected = pd.arrays.BooleanArray(values, mask)
+        tm.assert_extension_array_equal(result, expected)
+
+        # ensure we haven't mutated anything inplace
+        result[0] = pd.NA
+        tm.assert_extension_array_equal(
+            left, pd.array([0, 1, 2, None, None, None], dtype=dtype)
+        )
+        tm.assert_extension_array_equal(
+            right, pd.array([0, 1, None, 0, 1, None], dtype=dtype)
+        )
+
+    def test_compare_with_booleanarray(self, all_compare_operators, dtype):
+        op = self.get_op_from_name(all_compare_operators)
+
+        left = pd.array([True, False, None] * 3, dtype="boolean")
+        right = pd.array([0] * 3 + [1] * 3 + [None] * 3, dtype=dtype)
+        other = pd.array([False] * 3 + [True] * 3 + [None] * 3, dtype="boolean")
+
+        expected = op(left, other)
+        result = op(left, right)
+        tm.assert_extension_array_equal(result, expected)
+
+        # reversed op
+        expected = op(other, left)
+        result = op(right, left)
+        tm.assert_extension_array_equal(result, expected)
+
+    def test_compare_to_string(self, dtype):
+        # GH#28930
+        ser = pd.Series([1, None], dtype=dtype)
+        result = ser == "a"
+        expected = pd.Series([False, pd.NA], dtype="boolean")
+
+        self.assert_series_equal(result, expected)