use tm.assert_equal instead of parametrizing assert funcs (#22995)

* Use tm.assert_equal instead of parametrizing assert_func

* Extend assert_equal

* Use tm.assert_equal in more places

* typo fixup

Author: jbrockmendel

doc/source/contributing.rst

@@ -880,7 +880,7 @@ If your change involves checking that a warning is actually emitted, use
 
 .. code-block:: python
 
-   with tm.assert_prodcues_warning(FutureWarning):
+   with tm.assert_produces_warning(FutureWarning):
        df.some_operation()
 
 We prefer this to the ``pytest.warns`` context manager because ours checks that the warning's

pandas/tests/indexes/datetimes/test_tools.py

@@ -180,9 +180,9 @@ def test_to_datetime_format_weeks(self, cache):
         for s, format, dt in data:
             assert to_datetime(s, format=format, cache=cache) == dt
 
-    @pytest.mark.parametrize("box,const,assert_equal", [
-        [True, pd.Index, 'assert_index_equal'],
-        [False, np.array, 'assert_numpy_array_equal']])
+    @pytest.mark.parametrize("box,const", [
+        [True, pd.Index],
+        [False, np.array]])
     @pytest.mark.parametrize("fmt,dates,expected_dates", [
         ['%Y-%m-%d %H:%M:%S %Z',
          ['2010-01-01 12:00:00 UTC'] * 2,
@@ -215,12 +215,11 @@ def test_to_datetime_format_weeks(self, cache):
           pd.Timestamp('2010-01-01 12:00:00',
                        tzinfo=pytz.FixedOffset(0))]]])
     def test_to_datetime_parse_tzname_or_tzoffset(self, box, const,
-                                                  assert_equal, fmt,
-                                                  dates, expected_dates):
+                                                  fmt, dates, expected_dates):
         # GH 13486
         result = pd.to_datetime(dates, format=fmt, box=box)
         expected = const(expected_dates)
-        getattr(tm, assert_equal)(result, expected)
+        tm.assert_equal(result, expected)
 
         with pytest.raises(ValueError):
             pd.to_datetime(dates, format=fmt, box=box, utc=True)
@@ -1049,17 +1048,16 @@ def test_to_datetime_types(self, cache):
         # assert result == expected
 
     @pytest.mark.parametrize('cache', [True, False])
-    @pytest.mark.parametrize('box, klass, assert_method', [
-        [True, Index, 'assert_index_equal'],
-        [False, np.array, 'assert_numpy_array_equal']
+    @pytest.mark.parametrize('box, klass', [
+        [True, Index],
+        [False, np.array]
     ])
-    def test_to_datetime_unprocessable_input(self, cache, box, klass,
-                                             assert_method):
+    def test_to_datetime_unprocessable_input(self, cache, box, klass):
         # GH 4928
         # GH 21864
         result = to_datetime([1, '1'], errors='ignore', cache=cache, box=box)
         expected = klass(np.array([1, '1'], dtype='O'))
-        getattr(tm, assert_method)(result, expected)
+        tm.assert_equal(result, expected)
         pytest.raises(TypeError, to_datetime, [1, '1'], errors='raise',
                       cache=cache, box=box)
 

pandas/tests/scalar/test_nat.py

@@ -312,19 +312,16 @@ def test_nat_arithmetic_index():
     tm.assert_index_equal(NaT - tdi, tdi_nat)
 
 
-@pytest.mark.parametrize('box, assert_func', [
-    (TimedeltaIndex, tm.assert_index_equal),
-    (Series, tm.assert_series_equal)
-])
-def test_nat_arithmetic_td64_vector(box, assert_func):
+@pytest.mark.parametrize('box', [TimedeltaIndex, Series])
+def test_nat_arithmetic_td64_vector(box):
     # GH#19124
     vec = box(['1 day', '2 day'], dtype='timedelta64[ns]')
     box_nat = box([NaT, NaT], dtype='timedelta64[ns]')
 
-    assert_func(vec + NaT, box_nat)
-    assert_func(NaT + vec, box_nat)
-    assert_func(vec - NaT, box_nat)
-    assert_func(NaT - vec, box_nat)
+    tm.assert_equal(vec + NaT, box_nat)
+    tm.assert_equal(NaT + vec, box_nat)
+    tm.assert_equal(vec - NaT, box_nat)
+    tm.assert_equal(NaT - vec, box_nat)
 
 
 def test_nat_pinned_docstrings():

pandas/tests/tseries/offsets/test_offsets.py

@@ -2516,28 +2516,26 @@ def test_onOffset(self, case):
         dt, expected = case
         assert_onOffset(SemiMonthEnd(), dt, expected)
 
-    @pytest.mark.parametrize('klass,assert_func',
-                             [(Series, tm.assert_series_equal),
-                              (DatetimeIndex, tm.assert_index_equal)])
-    def test_vectorized_offset_addition(self, klass, assert_func):
+    @pytest.mark.parametrize('klass', [Series, DatetimeIndex])
+    def test_vectorized_offset_addition(self, klass):
         s = klass([Timestamp('2000-01-15 00:15:00', tz='US/Central'),
                    Timestamp('2000-02-15', tz='US/Central')], name='a')
 
         result = s + SemiMonthEnd()
         result2 = SemiMonthEnd() + s
         exp = klass([Timestamp('2000-01-31 00:15:00', tz='US/Central'),
                      Timestamp('2000-02-29', tz='US/Central')], name='a')
-        assert_func(result, exp)
-        assert_func(result2, exp)
+        tm.assert_equal(result, exp)
+        tm.assert_equal(result2, exp)
 
         s = klass([Timestamp('2000-01-01 00:15:00', tz='US/Central'),
                    Timestamp('2000-02-01', tz='US/Central')], name='a')
         result = s + SemiMonthEnd()
         result2 = SemiMonthEnd() + s
         exp = klass([Timestamp('2000-01-15 00:15:00', tz='US/Central'),
                      Timestamp('2000-02-15', tz='US/Central')], name='a')
-        assert_func(result, exp)
-        assert_func(result2, exp)
+        tm.assert_equal(result, exp)
+        tm.assert_equal(result2, exp)
 
 
 class TestSemiMonthBegin(Base):
@@ -2692,27 +2690,25 @@ def test_onOffset(self, case):
         dt, expected = case
         assert_onOffset(SemiMonthBegin(), dt, expected)
 
-    @pytest.mark.parametrize('klass,assert_func',
-                             [(Series, tm.assert_series_equal),
-                              (DatetimeIndex, tm.assert_index_equal)])
-    def test_vectorized_offset_addition(self, klass, assert_func):
+    @pytest.mark.parametrize('klass', [Series, DatetimeIndex])
+    def test_vectorized_offset_addition(self, klass):
         s = klass([Timestamp('2000-01-15 00:15:00', tz='US/Central'),
                    Timestamp('2000-02-15', tz='US/Central')], name='a')
         result = s + SemiMonthBegin()
         result2 = SemiMonthBegin() + s
         exp = klass([Timestamp('2000-02-01 00:15:00', tz='US/Central'),
                      Timestamp('2000-03-01', tz='US/Central')], name='a')
-        assert_func(result, exp)
-        assert_func(result2, exp)
+        tm.assert_equal(result, exp)
+        tm.assert_equal(result2, exp)
 
         s = klass([Timestamp('2000-01-01 00:15:00', tz='US/Central'),
                    Timestamp('2000-02-01', tz='US/Central')], name='a')
         result = s + SemiMonthBegin()
         result2 = SemiMonthBegin() + s
         exp = klass([Timestamp('2000-01-15 00:15:00', tz='US/Central'),
                      Timestamp('2000-02-15', tz='US/Central')], name='a')
-        assert_func(result, exp)
-        assert_func(result2, exp)
+        tm.assert_equal(result, exp)
+        tm.assert_equal(result2, exp)
 
 
 def test_Easter():

pandas/util/testing.py

@@ -1522,8 +1522,8 @@ def assert_equal(left, right, **kwargs):
 
     Parameters
     ----------
-    left : Index, Series, or DataFrame
-    right : Index, Series, or DataFrame
+    left : Index, Series, DataFrame, ExtensionArray, or np.ndarray
+    right : Index, Series, DataFrame, ExtensionArray, or np.ndarray
     **kwargs
     """
     if isinstance(left, pd.Index):
@@ -1532,6 +1532,10 @@ def assert_equal(left, right, **kwargs):
         assert_series_equal(left, right, **kwargs)
     elif isinstance(left, pd.DataFrame):
         assert_frame_equal(left, right, **kwargs)
+    elif isinstance(left, ExtensionArray):
+        assert_extension_array_equal(left, right, **kwargs)
+    elif isinstance(left, np.ndarray):
+        assert_numpy_array_equal(left, right, **kwargs)
     else:
         raise NotImplementedError(type(left))