diff --git a/pandas/computation/expr.py b/pandas/computation/expr.py index b6a1fcbec8339..73ccd11e966d9 100644 --- a/pandas/computation/expr.py +++ b/pandas/computation/expr.py @@ -589,6 +589,7 @@ def visitor(x, y): _python_not_supported = frozenset(['Dict', 'Call', 'BoolOp', 'In', 'NotIn']) _numexpr_supported_calls = frozenset(_reductions + _mathops) +_query_not_supported = frozenset(['Assign']) @disallow((_unsupported_nodes | _python_not_supported) - @@ -602,6 +603,17 @@ def __init__(self, env, engine, parser, super(PandasExprVisitor, self).__init__(env, engine, parser, preparser) +@disallow((_unsupported_nodes | _python_not_supported | _query_not_supported) - + (_boolop_nodes | frozenset(['BoolOp', 'Attribute', 'In', 'NotIn', + 'Tuple']))) +class PandasQueryExprVisitor(BaseExprVisitor): + + def __init__(self, env, engine, parser, + preparser=partial(_preparse, f=compose(_replace_locals, + _replace_booleans))): + super(PandasQueryExprVisitor, self).__init__(env, engine, parser, preparser) + + @disallow(_unsupported_nodes | _python_not_supported | frozenset(['Not'])) class PythonExprVisitor(BaseExprVisitor): @@ -659,4 +671,5 @@ def names(self): return frozenset(term.name for term in com.flatten(self.terms)) -_parsers = {'python': PythonExprVisitor, 'pandas': PandasExprVisitor} +_parsers = {'python': PythonExprVisitor, 'pandas': PandasExprVisitor, + 'pandas_query': PandasQueryExprVisitor} diff --git a/pandas/core/frame.py b/pandas/core/frame.py index bb192aeca5b6d..163a78d0a28a8 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -1935,6 +1935,7 @@ def query(self, expr, **kwargs): >>> df[df.a > df.b] # same result as the previous expression """ kwargs['level'] = kwargs.pop('level', 0) + 1 + kwargs['parser'] = kwargs.pop('parser', 'pandas_query') res = self.eval(expr, **kwargs) try: diff --git a/pandas/tests/test_frame.py b/pandas/tests/test_frame.py index cfa0ed1a11772..eda6d833f2b64 100644 --- a/pandas/tests/test_frame.py +++ b/pandas/tests/test_frame.py @@ -15210,7 +15210,7 @@ class TestDataFrameQueryPythonPandas(TestDataFrameQueryNumExprPandas): def setUpClass(cls): super(TestDataFrameQueryPythonPandas, cls).setUpClass() cls.engine = 'python' - cls.parser = 'pandas' + cls.parser = 'pandas_query' cls.frame = _frame.copy() def test_query_builtin(self): @@ -15225,6 +15225,16 @@ def test_query_builtin(self): result = df.query('sin > 5', engine=engine, parser=parser) tm.assert_frame_equal(expected, result) + def test_query_with_assign_statement(self): + df = DataFrame({'a': [1, 2, 3], 'b': ['a', 'b', 'c']}) + a_before = df['a'].copy() + self.assertRaisesRegexp( + NotImplementedError, "'Assign' nodes are not implemented", + df.query, 'a=1', engine=self.engine, parser=self.parser + ) + a_after = df['a'].copy() + assert_series_equal(a_before, a_after) + class TestDataFrameQueryPythonPython(TestDataFrameQueryNumExprPython):