TST/CLN: parametrize and clean test_expressions, test_nanops (pandas-dev#28553)

Author: jbrockmendel

pandas/tests/test_expressions.py

@@ -54,14 +54,12 @@ def run_arithmetic(self, df, other):
         operations = ["add", "sub", "mul", "mod", "truediv", "floordiv"]
         for test_flex in [True, False]:
             for arith in operations:
-
-                operator_name = arith
-
+                # TODO: share with run_binary
                 if test_flex:
                     op = lambda x, y: getattr(x, arith)(y)
                     op.__name__ = arith
                 else:
-                    op = getattr(operator, operator_name)
+                    op = getattr(operator, arith)
                 expr.set_use_numexpr(False)
                 expected = op(df, other)
                 expr.set_use_numexpr(True)
@@ -87,13 +85,14 @@ def run_binary(self, df, other):
         for test_flex in [True, False]:
             for arith in operations:
                 if test_flex:
-                    op = lambda x, y: getattr(df, arith)(y)
+                    op = lambda x, y: getattr(x, arith)(y)
                     op.__name__ = arith
                 else:
                     op = getattr(operator, arith)
                 expr.set_use_numexpr(False)
                 expected = op(df, other)
                 expr.set_use_numexpr(True)
+
                 expr.get_test_result()
                 result = op(df, other)
                 used_numexpr = expr.get_test_result()
@@ -167,29 +166,29 @@ def test_invalid(self):
         "opname,op_str",
         [("add", "+"), ("sub", "-"), ("mul", "*"), ("truediv", "/"), ("pow", "**")],
     )
-    def test_binary_ops(self, opname, op_str):
+    @pytest.mark.parametrize("left,right", [(_frame, _frame2), (_mixed, _mixed2)])
+    def test_binary_ops(self, opname, op_str, left, right):
         def testit():
 
-            for f, f2 in [(self.frame, self.frame2), (self.mixed, self.mixed2)]:
+            if opname == "pow":
+                # TODO: get this working
+                return
 
-                if opname == "pow":
-                    continue
+            op = getattr(operator, opname)
 
-                op = getattr(operator, opname)
+            result = expr._can_use_numexpr(op, op_str, left, left, "evaluate")
+            assert result != left._is_mixed_type
 
-                result = expr._can_use_numexpr(op, op_str, f, f, "evaluate")
-                assert result != f._is_mixed_type
+            result = expr.evaluate(op, op_str, left, left, use_numexpr=True)
+            expected = expr.evaluate(op, op_str, left, left, use_numexpr=False)
 
-                result = expr.evaluate(op, op_str, f, f, use_numexpr=True)
-                expected = expr.evaluate(op, op_str, f, f, use_numexpr=False)
+            if isinstance(result, DataFrame):
+                tm.assert_frame_equal(result, expected)
+            else:
+                tm.assert_numpy_array_equal(result, expected.values)
 
-                if isinstance(result, DataFrame):
-                    tm.assert_frame_equal(result, expected)
-                else:
-                    tm.assert_numpy_array_equal(result, expected.values)
-
-                result = expr._can_use_numexpr(op, op_str, f2, f2, "evaluate")
-                assert not result
+            result = expr._can_use_numexpr(op, op_str, right, right, "evaluate")
+            assert not result
 
         expr.set_use_numexpr(False)
         testit()
@@ -210,30 +209,26 @@ def testit():
             ("ne", "!="),
         ],
     )
-    def test_comparison_ops(self, opname, op_str):
+    @pytest.mark.parametrize("left,right", [(_frame, _frame2), (_mixed, _mixed2)])
+    def test_comparison_ops(self, opname, op_str, left, right):
         def testit():
-            for f, f2 in [(self.frame, self.frame2), (self.mixed, self.mixed2)]:
-
-                f11 = f
-                f12 = f + 1
+            f12 = left + 1
+            f22 = right + 1
 
-                f21 = f2
-                f22 = f2 + 1
+            op = getattr(operator, opname)
 
-                op = getattr(operator, opname)
+            result = expr._can_use_numexpr(op, op_str, left, f12, "evaluate")
+            assert result != left._is_mixed_type
 
-                result = expr._can_use_numexpr(op, op_str, f11, f12, "evaluate")
-                assert result != f11._is_mixed_type
+            result = expr.evaluate(op, op_str, left, f12, use_numexpr=True)
+            expected = expr.evaluate(op, op_str, left, f12, use_numexpr=False)
+            if isinstance(result, DataFrame):
+                tm.assert_frame_equal(result, expected)
+            else:
+                tm.assert_numpy_array_equal(result, expected.values)
 
-                result = expr.evaluate(op, op_str, f11, f12, use_numexpr=True)
-                expected = expr.evaluate(op, op_str, f11, f12, use_numexpr=False)
-                if isinstance(result, DataFrame):
-                    tm.assert_frame_equal(result, expected)
-                else:
-                    tm.assert_numpy_array_equal(result, expected.values)
-
-                result = expr._can_use_numexpr(op, op_str, f21, f22, "evaluate")
-                assert not result
+            result = expr._can_use_numexpr(op, op_str, right, f22, "evaluate")
+            assert not result
 
         expr.set_use_numexpr(False)
         testit()
@@ -244,15 +239,14 @@ def testit():
         testit()
 
     @pytest.mark.parametrize("cond", [True, False])
-    def test_where(self, cond):
+    @pytest.mark.parametrize("df", [_frame, _frame2, _mixed, _mixed2])
+    def test_where(self, cond, df):
         def testit():
-            for f in [self.frame, self.frame2, self.mixed, self.mixed2]:
-
-                c = np.empty(f.shape, dtype=np.bool_)
-                c.fill(cond)
-                result = expr.where(c, f.values, f.values + 1)
-                expected = np.where(c, f.values, f.values + 1)
-                tm.assert_numpy_array_equal(result, expected)
+            c = np.empty(df.shape, dtype=np.bool_)
+            c.fill(cond)
+            result = expr.where(c, df.values, df.values + 1)
+            expected = np.where(c, df.values, df.values + 1)
+            tm.assert_numpy_array_equal(result, expected)
 
         expr.set_use_numexpr(False)
         testit()
@@ -263,7 +257,7 @@ def testit():
         testit()
 
     @pytest.mark.parametrize(
-        "op_str,opname", list(zip(["/", "//", "**"], ["truediv", "floordiv", "pow"]))
+        "op_str,opname", [("/", "truediv"), ("//", "floordiv"), ("**", "pow")]
     )
     def test_bool_ops_raise_on_arithmetic(self, op_str, opname):
         df = DataFrame({"a": np.random.rand(10) > 0.5, "b": np.random.rand(10) > 0.5})
@@ -291,7 +285,7 @@ def test_bool_ops_raise_on_arithmetic(self, op_str, opname):
             f(df, True)
 
     @pytest.mark.parametrize(
-        "op_str,opname", list(zip(["+", "*", "-"], ["add", "mul", "sub"]))
+        "op_str,opname", [("+", "add"), ("*", "mul"), ("-", "sub")]
     )
     def test_bool_ops_warn_on_arithmetic(self, op_str, opname):
         n = 10