REF: Base class for all extension tests (pandas-dev#19863)

Author: TomAugspurger

ci/lint.sh

@@ -111,6 +111,15 @@ if [ "$LINT" ]; then
         RET=1
     fi
 
+    # Check for the following code in the extension array base tests
+    # tm.assert_frame_equal
+    # tm.assert_series_equal
+    grep -r -E --include '*.py' --exclude base.py 'tm.assert_(series|frame)_equal' pandas/tests/extension/base
+
+    if [ $? = "0" ]; then
+        RET=1
+    fi
+
     echo "Check for invalid testing DONE"
 
     # Check for imports from pandas.core.common instead

pandas/tests/extension/base/__init__.py

@@ -31,6 +31,14 @@ class TestMyDtype(BaseDtypeTests):
 Your class ``TestDtype`` will inherit all the tests defined on
 ``BaseDtypeTests``. pytest's fixture discover will supply your ``dtype``
 wherever the test requires it. You're free to implement additional tests.
+
+All the tests in these modules use ``self.assert_frame_equal`` or
+``self.assert_series_equal`` for dataframe or series comparisons. By default,
+they use the usual ``pandas.testing.assert_frame_equal`` and
+``pandas.testing.assert_series_equal``. You can override the checks used
+by defining the staticmethods ``assert_frame_equal`` and
+``assert_series_equal`` on your base test class.
+
 """
 from .casting import BaseCastingTests  # noqa
 from .constructors import BaseConstructorsTests  # noqa

pandas/tests/extension/base/base.py

@@ -0,0 +1,6 @@
+import pandas.util.testing as tm
+
+
+class BaseExtensionTests(object):
+    assert_series_equal = staticmethod(tm.assert_series_equal)
+    assert_frame_equal = staticmethod(tm.assert_frame_equal)

pandas/tests/extension/base/casting.py

@@ -1,8 +1,10 @@
 import pandas as pd
 from pandas.core.internals import ObjectBlock
 
+from .base import BaseExtensionTests
 
-class BaseCastingTests(object):
+
+class BaseCastingTests(BaseExtensionTests):
     """Casting to and from ExtensionDtypes"""
 
     def test_astype_object_series(self, all_data):

pandas/tests/extension/base/constructors.py

@@ -4,8 +4,10 @@
 import pandas.util.testing as tm
 from pandas.core.internals import ExtensionBlock
 
+from .base import BaseExtensionTests
 
-class BaseConstructorsTests(object):
+
+class BaseConstructorsTests(BaseExtensionTests):
 
     def test_series_constructor(self, data):
         result = pd.Series(data)

pandas/tests/extension/base/dtype.py

@@ -1,8 +1,10 @@
 import numpy as np
 import pandas as pd
 
+from .base import BaseExtensionTests
 
-class BaseDtypeTests(object):
+
+class BaseDtypeTests(BaseExtensionTests):
     """Base class for ExtensionDtype classes"""
 
     def test_name(self, dtype):

pandas/tests/extension/base/getitem.py

@@ -1,73 +1,74 @@
 import numpy as np
 
 import pandas as pd
-import pandas.util.testing as tm
 
+from .base import BaseExtensionTests
 
-class BaseGetitemTests(object):
+
+class BaseGetitemTests(BaseExtensionTests):
     """Tests for ExtensionArray.__getitem__."""
 
     def test_iloc_series(self, data):
         ser = pd.Series(data)
         result = ser.iloc[:4]
         expected = pd.Series(data[:4])
-        tm.assert_series_equal(result, expected)
+        self.assert_series_equal(result, expected)
 
         result = ser.iloc[[0, 1, 2, 3]]
-        tm.assert_series_equal(result, expected)
+        self.assert_series_equal(result, expected)
 
     def test_iloc_frame(self, data):
         df = pd.DataFrame({"A": data, 'B': np.arange(len(data))})
         expected = pd.DataFrame({"A": data[:4]})
 
         # slice -> frame
         result = df.iloc[:4, [0]]
-        tm.assert_frame_equal(result, expected)
+        self.assert_frame_equal(result, expected)
 
         # sequence -> frame
         result = df.iloc[[0, 1, 2, 3], [0]]
-        tm.assert_frame_equal(result, expected)
+        self.assert_frame_equal(result, expected)
 
         expected = pd.Series(data[:4], name='A')
 
         # slice -> series
         result = df.iloc[:4, 0]
-        tm.assert_series_equal(result, expected)
+        self.assert_series_equal(result, expected)
 
         # sequence -> series
         result = df.iloc[:4, 0]
-        tm.assert_series_equal(result, expected)
+        self.assert_series_equal(result, expected)
 
     def test_loc_series(self, data):
         ser = pd.Series(data)
         result = ser.loc[:3]
         expected = pd.Series(data[:4])
-        tm.assert_series_equal(result, expected)
+        self.assert_series_equal(result, expected)
 
         result = ser.loc[[0, 1, 2, 3]]
-        tm.assert_series_equal(result, expected)
+        self.assert_series_equal(result, expected)
 
     def test_loc_frame(self, data):
         df = pd.DataFrame({"A": data, 'B': np.arange(len(data))})
         expected = pd.DataFrame({"A": data[:4]})
 
         # slice -> frame
         result = df.loc[:3, ['A']]
-        tm.assert_frame_equal(result, expected)
+        self.assert_frame_equal(result, expected)
 
         # sequence -> frame
         result = df.loc[[0, 1, 2, 3], ['A']]
-        tm.assert_frame_equal(result, expected)
+        self.assert_frame_equal(result, expected)
 
         expected = pd.Series(data[:4], name='A')
 
         # slice -> series
         result = df.loc[:3, 'A']
-        tm.assert_series_equal(result, expected)
+        self.assert_series_equal(result, expected)
 
         # sequence -> series
         result = df.loc[:3, 'A']
-        tm.assert_series_equal(result, expected)
+        self.assert_series_equal(result, expected)
 
     def test_getitem_scalar(self, data):
         result = data[0]

pandas/tests/extension/base/interface.py

@@ -5,8 +5,10 @@
 from pandas.core.dtypes.common import is_extension_array_dtype
 from pandas.core.dtypes.dtypes import ExtensionDtype
 
+from .base import BaseExtensionTests
 
-class BaseInterfaceTests(object):
+
+class BaseInterfaceTests(BaseExtensionTests):
     """Tests that the basic interface is satisfied."""
     # ------------------------------------------------------------------------
     # Interface

pandas/tests/extension/base/methods.py

@@ -2,10 +2,11 @@
 import numpy as np
 
 import pandas as pd
-import pandas.util.testing as tm
 
+from .base import BaseExtensionTests
 
-class BaseMethodsTests(object):
+
+class BaseMethodsTests(BaseExtensionTests):
     """Various Series and DataFrame methods."""
 
     @pytest.mark.parametrize('dropna', [True, False])
@@ -19,13 +20,13 @@ def test_value_counts(self, all_data, dropna):
         result = pd.Series(all_data).value_counts(dropna=dropna).sort_index()
         expected = pd.Series(other).value_counts(dropna=dropna).sort_index()
 
-        tm.assert_series_equal(result, expected)
+        self.assert_series_equal(result, expected)
 
     def test_count(self, data_missing):
         df = pd.DataFrame({"A": data_missing})
         result = df.count(axis='columns')
         expected = pd.Series([0, 1])
-        tm.assert_series_equal(result, expected)
+        self.assert_series_equal(result, expected)
 
     def test_apply_simple_series(self, data):
         result = pd.Series(data).apply(id)

pandas/tests/extension/base/missing.py

@@ -3,8 +3,10 @@
 import pandas as pd
 import pandas.util.testing as tm
 
+from .base import BaseExtensionTests
 
-class BaseMissingTests(object):
+
+class BaseMissingTests(BaseExtensionTests):
     def test_isna(self, data_missing):
         if data_missing._can_hold_na:
             expected = np.array([True, False])
@@ -16,30 +18,30 @@ def test_isna(self, data_missing):
 
         result = pd.Series(data_missing).isna()
         expected = pd.Series(expected)
-        tm.assert_series_equal(result, expected)
+        self.assert_series_equal(result, expected)
 
     def test_dropna_series(self, data_missing):
         ser = pd.Series(data_missing)
         result = ser.dropna()
         expected = ser.iloc[[1]]
-        tm.assert_series_equal(result, expected)
+        self.assert_series_equal(result, expected)
 
     def test_dropna_frame(self, data_missing):
         df = pd.DataFrame({"A": data_missing})
 
         # defaults
         result = df.dropna()
         expected = df.iloc[[1]]
-        tm.assert_frame_equal(result, expected)
+        self.assert_frame_equal(result, expected)
 
         # axis = 1
         result = df.dropna(axis='columns')
         expected = pd.DataFrame(index=[0, 1])
-        tm.assert_frame_equal(result, expected)
+        self.assert_frame_equal(result, expected)
 
         # multiple
         df = pd.DataFrame({"A": data_missing,
                            "B": [1, np.nan]})
         result = df.dropna()
         expected = df.iloc[:0]
-        tm.assert_frame_equal(result, expected)
+        self.assert_frame_equal(result, expected)

pandas/tests/extension/base/reshaping.py

@@ -1,11 +1,12 @@
 import pytest
 
 import pandas as pd
-import pandas.util.testing as tm
 from pandas.core.internals import ExtensionBlock
 
+from .base import BaseExtensionTests
 
-class BaseReshapingTests(object):
+
+class BaseReshapingTests(BaseExtensionTests):
     """Tests for reshaping and concatenation."""
     @pytest.mark.parametrize('in_frame', [True, False])
     def test_concat(self, data, in_frame):
@@ -32,8 +33,8 @@ def test_align(self, data, na_value):
         # Assumes that the ctor can take a list of scalars of the type
         e1 = pd.Series(type(data)(list(a) + [na_value]))
         e2 = pd.Series(type(data)([na_value] + list(b)))
-        tm.assert_series_equal(r1, e1)
-        tm.assert_series_equal(r2, e2)
+        self.assert_series_equal(r1, e1)
+        self.assert_series_equal(r2, e2)
 
     def test_align_frame(self, data, na_value):
         a = data[:3]
@@ -45,17 +46,17 @@ def test_align_frame(self, data, na_value):
         # Assumes that the ctor can take a list of scalars of the type
         e1 = pd.DataFrame({'A': type(data)(list(a) + [na_value])})
         e2 = pd.DataFrame({'A': type(data)([na_value] + list(b))})
-        tm.assert_frame_equal(r1, e1)
-        tm.assert_frame_equal(r2, e2)
+        self.assert_frame_equal(r1, e1)
+        self.assert_frame_equal(r2, e2)
 
     def test_set_frame_expand_regular_with_extension(self, data):
         df = pd.DataFrame({"A": [1] * len(data)})
         df['B'] = data
         expected = pd.DataFrame({"A": [1] * len(data), "B": data})
-        tm.assert_frame_equal(df, expected)
+        self.assert_frame_equal(df, expected)
 
     def test_set_frame_expand_extension_with_regular(self, data):
         df = pd.DataFrame({'A': data})
         df['B'] = [1] * len(data)
         expected = pd.DataFrame({"A": data, "B": [1] * len(data)})
-        tm.assert_frame_equal(df, expected)
+        self.assert_frame_equal(df, expected)