From ccf0b5307de2d923ad1cf45507d115a481ca0cc6 Mon Sep 17 00:00:00 2001 From: "alexprincel@gmail.com" Date: Mon, 22 Mar 2021 01:20:42 -0400 Subject: [PATCH] BUG: Fix query doesn't support equals or not equals with Python parser (#40436) --- doc/source/whatsnew/v1.3.0.rst | 1 + pandas/core/computation/expr.py | 2 -- pandas/tests/frame/test_query_eval.py | 7 +++++++ 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/doc/source/whatsnew/v1.3.0.rst b/doc/source/whatsnew/v1.3.0.rst index 8deeb3cfae1d3..52751dc607fcd 100644 --- a/doc/source/whatsnew/v1.3.0.rst +++ b/doc/source/whatsnew/v1.3.0.rst @@ -494,6 +494,7 @@ Other - Bug in :meth:`Styler.background_gradient` where text-color was not determined correctly (:issue:`39888`) - Bug in :meth:`DataFrame.equals`, :meth:`Series.equals`, :meth:`Index.equals` with object-dtype containing ``np.datetime64("NaT")`` or ``np.timedelta64("NaT")`` (:issue:`39650`) - Bug in :func:`pandas.util.show_versions` where console JSON output was not proper JSON (:issue:`39701`) +- Bug in :method:`DataFrame.query`, where "python" ``parser`` would raise NotImplementedError for == and != (:issue:`40436`) .. --------------------------------------------------------------------------- diff --git a/pandas/core/computation/expr.py b/pandas/core/computation/expr.py index 02660539f4981..aee5c3442780a 100644 --- a/pandas/core/computation/expr.py +++ b/pandas/core/computation/expr.py @@ -384,8 +384,6 @@ class BaseExprVisitor(ast.NodeVisitor): unary_op_nodes_map = {k: v for k, v in zip(unary_ops, unary_op_nodes)} rewrite_map = { - ast.Eq: ast.In, - ast.NotEq: ast.NotIn, ast.In: ast.In, ast.NotIn: ast.NotIn, } diff --git a/pandas/tests/frame/test_query_eval.py b/pandas/tests/frame/test_query_eval.py index fdbf8a93ddddf..785ad1001f537 100644 --- a/pandas/tests/frame/test_query_eval.py +++ b/pandas/tests/frame/test_query_eval.py @@ -548,6 +548,13 @@ def test_query(self): df[df.a + df.b > df.b * df.c], ) + def test_query_equality(self): + engine, parser = self.engine, self.parser + df = DataFrame({"x": np.random.choice(["A", "B", "C"], size=20)}) + tm.assert_frame_equal( + df.query("x == 'A'", engine=engine, parser=parser), df[df.x == "A"] + ) + def test_query_index_with_name(self): engine, parser = self.engine, self.parser df = DataFrame(