diff --git a/pandas/tests/frame/test_query_eval.py b/pandas/tests/frame/test_query_eval.py index 703e05998e93c..bf9eeb532b43b 100644 --- a/pandas/tests/frame/test_query_eval.py +++ b/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: