TST: parametrize eval tests (#31901)

Author: jbrockmendel

pandas/tests/frame/test_query_eval.py

@@ -78,45 +78,48 @@ def test_query_numexpr(self):
 
 
 class TestDataFrameEval:
-    def test_ops(self):
+
+    # smaller hits python, larger hits numexpr
+    @pytest.mark.parametrize("n", [4, 4000])
+    @pytest.mark.parametrize(
+        "op_str,op,rop",
+        [
+            ("+", "__add__", "__radd__"),
+            ("-", "__sub__", "__rsub__"),
+            ("*", "__mul__", "__rmul__"),
+            ("/", "__truediv__", "__rtruediv__"),
+        ],
+    )
+    def test_ops(self, op_str, op, rop, n):
 
         # tst ops and reversed ops in evaluation
         # GH7198
 
-        # smaller hits python, larger hits numexpr
-        for n in [4, 4000]:
-
-            df = DataFrame(1, index=range(n), columns=list("abcd"))
-            df.iloc[0] = 2
-            m = df.mean()
+        df = DataFrame(1, index=range(n), columns=list("abcd"))
+        df.iloc[0] = 2
+        m = df.mean()
 
-            for op_str, op, rop in [
-                ("+", "__add__", "__radd__"),
-                ("-", "__sub__", "__rsub__"),
-                ("*", "__mul__", "__rmul__"),
-                ("/", "__truediv__", "__rtruediv__"),
-            ]:
-
-                base = DataFrame(  # noqa
-                    np.tile(m.values, n).reshape(n, -1), columns=list("abcd")
-                )
+        base = DataFrame(  # noqa
+            np.tile(m.values, n).reshape(n, -1), columns=list("abcd")
+        )
 
-                expected = eval("base{op}df".format(op=op_str))
+        expected = eval(f"base {op_str} df")
 
-                # ops as strings
-                result = eval("m{op}df".format(op=op_str))
-                tm.assert_frame_equal(result, expected)
+        # ops as strings
+        result = eval(f"m {op_str} df")
+        tm.assert_frame_equal(result, expected)
 
-                # these are commutative
-                if op in ["+", "*"]:
-                    result = getattr(df, op)(m)
-                    tm.assert_frame_equal(result, expected)
+        # these are commutative
+        if op in ["+", "*"]:
+            result = getattr(df, op)(m)
+            tm.assert_frame_equal(result, expected)
 
-                # these are not
-                elif op in ["-", "/"]:
-                    result = getattr(df, rop)(m)
-                    tm.assert_frame_equal(result, expected)
+        # these are not
+        elif op in ["-", "/"]:
+            result = getattr(df, rop)(m)
+            tm.assert_frame_equal(result, expected)
 
+    def test_dataframe_sub_numexpr_path(self):
         # GH7192: Note we need a large number of rows to ensure this
         #  goes through the numexpr path
         df = DataFrame(dict(A=np.random.randn(25000)))
@@ -451,9 +454,7 @@ def test_date_query_with_non_date(self):
 
         for op in ["<", ">", "<=", ">="]:
             with pytest.raises(TypeError):
-                df.query(
-                    "dates {op} nondate".format(op=op), parser=parser, engine=engine
-                )
+                df.query(f"dates {op} nondate", parser=parser, engine=engine)
 
     def test_query_syntax_error(self):
         engine, parser = self.engine, self.parser
@@ -687,10 +688,9 @@ def test_inf(self):
         n = 10
         df = DataFrame({"a": np.random.rand(n), "b": np.random.rand(n)})
         df.loc[::2, 0] = np.inf
-        ops = "==", "!="
-        d = dict(zip(ops, (operator.eq, operator.ne)))
+        d = {"==": operator.eq, "!=": operator.ne}
         for op, f in d.items():
-            q = "a {op} inf".format(op=op)
+            q = f"a {op} inf"
             expected = df[f(df.a, np.inf)]
             result = df.query(q, engine=self.engine, parser=self.parser)
             tm.assert_frame_equal(result, expected)
@@ -854,7 +854,7 @@ def test_str_query_method(self, parser, engine):
             ops = 2 * ([eq] + [ne])
 
             for lhs, op, rhs in zip(lhs, ops, rhs):
-                ex = "{lhs} {op} {rhs}".format(lhs=lhs, op=op, rhs=rhs)
+                ex = f"{lhs} {op} {rhs}"
                 msg = r"'(Not)?In' nodes are not implemented"
                 with pytest.raises(NotImplementedError, match=msg):
                     df.query(
@@ -895,7 +895,7 @@ def test_str_list_query_method(self, parser, engine):
             ops = 2 * ([eq] + [ne])
 
             for lhs, op, rhs in zip(lhs, ops, rhs):
-                ex = "{lhs} {op} {rhs}".format(lhs=lhs, op=op, rhs=rhs)
+                ex = f"{lhs} {op} {rhs}"
                 with pytest.raises(NotImplementedError):
                     df.query(ex, engine=engine, parser=parser)
         else:
@@ -1042,7 +1042,7 @@ def test_invalid_type_for_operator_raises(self, parser, engine, op):
         msg = r"unsupported operand type\(s\) for .+: '.+' and '.+'"
 
         with pytest.raises(TypeError, match=msg):
-            df.eval("a {0} b".format(op), engine=engine, parser=parser)
+            df.eval(f"a {op} b", engine=engine, parser=parser)
 
 
 class TestDataFrameQueryBacktickQuoting: