From 9e1add75f771904ac124bbd0e313a77febe62ce4 Mon Sep 17 00:00:00 2001 From: Alessio Izzo Date: Tue, 30 May 2023 16:59:24 +0200 Subject: [PATCH] Fix warlus operator behavior when called by a function (#11041) In #10758 we introduced the support for the use of the walrus operator in the test cases. There was a case which was not handled that caused a bug report #11028. This PR aims to fix the issue and also to improve how the walrus operator is handled in the AssertionRewriter class. Closes #11028 --- changelog/11028.bugfix.rst | 1 + src/_pytest/assertion/rewrite.py | 25 +++++++-- testing/test_assertrewrite.py | 90 ++++++++++++++++++++++++++++++++ 3 files changed, 113 insertions(+), 3 deletions(-) create mode 100644 changelog/11028.bugfix.rst diff --git a/changelog/11028.bugfix.rst b/changelog/11028.bugfix.rst new file mode 100644 index 000000000..9efc04ba6 --- /dev/null +++ b/changelog/11028.bugfix.rst @@ -0,0 +1 @@ +Fixed bug in assertion rewriting where a variable assigned with the walrus operator could not be used later in a function call. diff --git a/src/_pytest/assertion/rewrite.py b/src/_pytest/assertion/rewrite.py index 8b1823470..ab8169da2 100644 --- a/src/_pytest/assertion/rewrite.py +++ b/src/_pytest/assertion/rewrite.py @@ -996,7 +996,9 @@ class AssertionRewriter(ast.NodeVisitor): ] ): pytest_temp = self.variable() - self.variables_overwrite[v.left.target.id] = pytest_temp + self.variables_overwrite[ + v.left.target.id + ] = v.left # type:ignore[assignment] v.left.target.id = pytest_temp self.push_format_context() res, expl = self.visit(v) @@ -1037,10 +1039,19 @@ class AssertionRewriter(ast.NodeVisitor): new_args = [] new_kwargs = [] for arg in call.args: + if isinstance(arg, ast.Name) and arg.id in self.variables_overwrite: + arg = self.variables_overwrite[arg.id] # type:ignore[assignment] res, expl = self.visit(arg) arg_expls.append(expl) new_args.append(res) for keyword in call.keywords: + if ( + isinstance(keyword.value, ast.Name) + and keyword.value.id in self.variables_overwrite + ): + keyword.value = self.variables_overwrite[ + keyword.value.id + ] # type:ignore[assignment] res, expl = self.visit(keyword.value) new_kwargs.append(ast.keyword(keyword.arg, res)) if keyword.arg: @@ -1075,7 +1086,13 @@ class AssertionRewriter(ast.NodeVisitor): self.push_format_context() # We first check if we have overwritten a variable in the previous assert if isinstance(comp.left, ast.Name) and comp.left.id in self.variables_overwrite: - comp.left.id = self.variables_overwrite[comp.left.id] + comp.left = self.variables_overwrite[ + comp.left.id + ] # type:ignore[assignment] + if isinstance(comp.left, namedExpr): + self.variables_overwrite[ + comp.left.target.id + ] = comp.left # type:ignore[assignment] left_res, left_expl = self.visit(comp.left) if isinstance(comp.left, (ast.Compare, ast.BoolOp)): left_expl = f"({left_expl})" @@ -1093,7 +1110,9 @@ class AssertionRewriter(ast.NodeVisitor): and next_operand.target.id == left_res.id ): next_operand.target.id = self.variable() - self.variables_overwrite[left_res.id] = next_operand.target.id + self.variables_overwrite[ + left_res.id + ] = next_operand # type:ignore[assignment] next_res, next_expl = self.visit(next_operand) if isinstance(next_operand, (ast.Compare, ast.BoolOp)): next_expl = f"({next_expl})" diff --git a/testing/test_assertrewrite.py b/testing/test_assertrewrite.py index 8d9441403..245241af2 100644 --- a/testing/test_assertrewrite.py +++ b/testing/test_assertrewrite.py @@ -1436,6 +1436,96 @@ class TestIssue10743: assert result.ret == 0 +@pytest.mark.skipif( + sys.version_info < (3, 8), reason="walrus operator not available in py<38" +) +class TestIssue11028: + def test_assertion_walrus_operator_in_operand(self, pytester: Pytester) -> None: + pytester.makepyfile( + """ + def test_in_string(): + assert (obj := "foo") in obj + """ + ) + result = pytester.runpytest() + assert result.ret == 0 + + def test_assertion_walrus_operator_in_operand_json_dumps( + self, pytester: Pytester + ) -> None: + pytester.makepyfile( + """ + import json + + def test_json_encoder(): + assert (obj := "foo") in json.dumps(obj) + """ + ) + result = pytester.runpytest() + assert result.ret == 0 + + def test_assertion_walrus_operator_equals_operand_function( + self, pytester: Pytester + ) -> None: + pytester.makepyfile( + """ + def f(a): + return a + + def test_call_other_function_arg(): + assert (obj := "foo") == f(obj) + """ + ) + result = pytester.runpytest() + assert result.ret == 0 + + def test_assertion_walrus_operator_equals_operand_function_keyword_arg( + self, pytester: Pytester + ) -> None: + pytester.makepyfile( + """ + def f(a='test'): + return a + + def test_call_other_function_k_arg(): + assert (obj := "foo") == f(a=obj) + """ + ) + result = pytester.runpytest() + assert result.ret == 0 + + def test_assertion_walrus_operator_equals_operand_function_arg_as_function( + self, pytester: Pytester + ) -> None: + pytester.makepyfile( + """ + def f(a='test'): + return a + + def test_function_of_function(): + assert (obj := "foo") == f(f(obj)) + """ + ) + result = pytester.runpytest() + assert result.ret == 0 + + def test_assertion_walrus_operator_gt_operand_function( + self, pytester: Pytester + ) -> None: + pytester.makepyfile( + """ + def add_one(a): + return a + 1 + + def test_gt(): + assert (obj := 4) > add_one(obj) + """ + ) + result = pytester.runpytest() + assert result.ret == 1 + result.stdout.fnmatch_lines(["*assert 4 > 5", "*where 5 = add_one(4)"]) + + @pytest.mark.skipif( sys.maxsize <= (2**31 - 1), reason="Causes OverflowError on 32bit systems" )