Skip to content

Commit

Permalink
Fix assert rewriting with assignment expressions
Browse files Browse the repository at this point in the history
  • Loading branch information
cdce8p committed Sep 8, 2023
1 parent 0a06db0 commit a9c6ac1
Show file tree
Hide file tree
Showing 4 changed files with 54 additions and 14 deletions.
1 change: 1 addition & 0 deletions AUTHORS
Original file line number Diff line number Diff line change
Expand Up @@ -235,6 +235,7 @@ Maho
Maik Figura
Mandeep Bhutani
Manuel Krebber
Marc Mueller
Marc Schlaich
Marcelo Duarte Trevisani
Marcin Bachry
Expand Down
1 change: 1 addition & 0 deletions changelog/11239.bugfix.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Fixed ``:=`` in asserts impacting unrelated test cases.
49 changes: 35 additions & 14 deletions src/_pytest/assertion/rewrite.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import sys
import tokenize
import types
from collections import defaultdict
from pathlib import Path
from pathlib import PurePath
from typing import Callable
Expand Down Expand Up @@ -52,6 +53,8 @@
PYC_EXT = ".py" + (__debug__ and "c" or "o")
PYC_TAIL = "." + PYTEST_TAG + PYC_EXT

_SCOPE_END_MARKER = object()


class AssertionRewritingHook(importlib.abc.MetaPathFinder, importlib.abc.Loader):
"""PEP302/PEP451 import hook which rewrites asserts."""
Expand Down Expand Up @@ -634,6 +637,8 @@ class AssertionRewriter(ast.NodeVisitor):
.push_format_context() and .pop_format_context() which allows
to build another %-formatted string while already building one.
:scope: A tuple containing the current scope used for variables_overwrite.
:variables_overwrite: A dict filled with references to variables
that change value within an assert. This happens when a variable is
reassigned with the walrus operator
Expand All @@ -655,7 +660,10 @@ def __init__(
else:
self.enable_assertion_pass_hook = False
self.source = source
self.variables_overwrite: Dict[str, str] = {}
self.scope: tuple[ast.AST, ...] = ()
self.variables_overwrite: defaultdict[
tuple[ast.AST, ...], Dict[str, str]
] = defaultdict(dict)

def run(self, mod: ast.Module) -> None:
"""Find all assert statements in *mod* and rewrite them."""
Expand Down Expand Up @@ -719,9 +727,17 @@ def run(self, mod: ast.Module) -> None:
mod.body[pos:pos] = imports

# Collect asserts.
nodes: List[ast.AST] = [mod]
self.scope = (mod,)
nodes: List[Union[ast.AST, object]] = [mod]
while nodes:
node = nodes.pop()
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
self.scope = tuple((*self.scope, node))
nodes.append(_SCOPE_END_MARKER)
if node == _SCOPE_END_MARKER:
self.scope = self.scope[:-1]
continue
assert isinstance(node, ast.AST)
for name, field in ast.iter_fields(node):
if isinstance(field, list):
new: List[ast.AST] = []
Expand Down Expand Up @@ -992,7 +1008,7 @@ def visit_BoolOp(self, boolop: ast.BoolOp) -> Tuple[ast.Name, str]:
]
):
pytest_temp = self.variable()
self.variables_overwrite[
self.variables_overwrite[self.scope][

Check warning on line 1011 in src/_pytest/assertion/rewrite.py

View check run for this annotation

Codecov / codecov/patch

src/_pytest/assertion/rewrite.py#L1011

Added line #L1011 was not covered by tests
v.left.target.id
] = v.left # type:ignore[assignment]
v.left.target.id = pytest_temp
Expand Down Expand Up @@ -1035,17 +1051,20 @@ def visit_Call(self, call: ast.Call) -> Tuple[ast.Name, str]:
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]
if isinstance(arg, ast.Name) and arg.id in self.variables_overwrite.get(
self.scope, {}
):
arg = self.variables_overwrite[self.scope][

Check warning on line 1057 in src/_pytest/assertion/rewrite.py

View check run for this annotation

Codecov / codecov/patch

src/_pytest/assertion/rewrite.py#L1057

Added line #L1057 was not covered by tests
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[
if isinstance(
keyword.value, ast.Name
) and keyword.value.id in self.variables_overwrite.get(self.scope, {}):
keyword.value = self.variables_overwrite[self.scope][

Check warning on line 1067 in src/_pytest/assertion/rewrite.py

View check run for this annotation

Codecov / codecov/patch

src/_pytest/assertion/rewrite.py#L1067

Added line #L1067 was not covered by tests
keyword.value.id
] # type:ignore[assignment]
res, expl = self.visit(keyword.value)
Expand Down Expand Up @@ -1081,12 +1100,14 @@ def visit_Attribute(self, attr: ast.Attribute) -> Tuple[ast.Name, str]:
def visit_Compare(self, comp: ast.Compare) -> Tuple[ast.expr, str]:
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 = self.variables_overwrite[
if isinstance(
comp.left, ast.Name
) and comp.left.id in self.variables_overwrite.get(self.scope, {}):
comp.left = self.variables_overwrite[self.scope][

Check warning on line 1106 in src/_pytest/assertion/rewrite.py

View check run for this annotation

Codecov / codecov/patch

src/_pytest/assertion/rewrite.py#L1106

Added line #L1106 was not covered by tests
comp.left.id
] # type:ignore[assignment]
if isinstance(comp.left, ast.NamedExpr):
self.variables_overwrite[
self.variables_overwrite[self.scope][

Check warning on line 1110 in src/_pytest/assertion/rewrite.py

View check run for this annotation

Codecov / codecov/patch

src/_pytest/assertion/rewrite.py#L1110

Added line #L1110 was not covered by tests
comp.left.target.id
] = comp.left # type:ignore[assignment]
left_res, left_expl = self.visit(comp.left)
Expand All @@ -1106,7 +1127,7 @@ def visit_Compare(self, comp: ast.Compare) -> Tuple[ast.expr, str]:
and next_operand.target.id == left_res.id
):
next_operand.target.id = self.variable()
self.variables_overwrite[
self.variables_overwrite[self.scope][

Check warning on line 1130 in src/_pytest/assertion/rewrite.py

View check run for this annotation

Codecov / codecov/patch

src/_pytest/assertion/rewrite.py#L1130

Added line #L1130 was not covered by tests
left_res.id
] = next_operand # type:ignore[assignment]
next_res, next_expl = self.visit(next_operand)
Expand Down
17 changes: 17 additions & 0 deletions testing/test_assertrewrite.py
Original file line number Diff line number Diff line change
Expand Up @@ -1543,6 +1543,23 @@ def test_gt():
result.stdout.fnmatch_lines(["*assert 4 > 5", "*where 5 = add_one(4)"])


class TestIssue11239:
def test_assertion_walrus_different_test_cases(self, pytester: Pytester) -> None:
pytester.makepyfile(
"""
def test_1():
state = {"x": 2}.get("x")
assert state is not None
def test_2():
db = {"x": 2}
assert (state := db.get("x")) is not None
"""
)
result = pytester.runpytest()
assert result.ret == 0


@pytest.mark.skipif(
sys.maxsize <= (2**31 - 1), reason="Causes OverflowError on 32bit systems"
)
Expand Down

0 comments on commit a9c6ac1

Please sign in to comment.