Fix assert rewriting with assignment expressions

[cdce8p]

Fixes #11239

Track the rudimentary scope where := are used to prevent replacing variables in other test cases. AFAICT it's not necessary to be precise and track every scope change as only assert nodes are visited anyway and those are not usually used in comprehensions, for example.

[nicoddemus]

Thanks @cdce8p, appreciate the contribution!

Please take a look at my comments. 👍

[nicoddemus]

Minor, but how about we make a specific subclass from ast.AST, so the typing for scope: tuple[ast.AST, ...] = () is consistent?

Suggested change

	_SCOPE_END_MARKER = object()
	class ScopeEnd(ast.AST):
	"""
	(some docs here).
	"""

And then instead of checking if node == _SCOPE_END_MARKER:. we can if isinstance(node, ScopeEnd).

[cdce8p]

The self.scope typing is correct with tuple[ast.AST, ...] already. _SCOE_END_MARKER is only added to the nodes list never to self.scope.

[nicoddemus]

I don't quite understand the role of _SCOPE_END_MARKER... could you explain it? Or perhaps if you follow my suggestion of using an AST subclass, add that explanation to the docstring of the subclass.

[cdce8p]

Although the AssertionRewriter inherits from NodeVisitor, not all nodes are in fact visited - only assert and their child nodes. That makes it a bit difficult to detect when a scope change happens. ATM pytest only iterates over all nodes recursively and copies them if they don't need to be rewritten. To do that, child nodes are added to nodes and then popped in the while loop. I added _SCOPE_END_MARKER to know when all child nodes e.g. for a FunctionDef have been visited.

An example

1. At the beginning, the nodes list contains maybe a few imports and say two FunctionDef nodes for test_1 and test_2.
2. In the while loop, the FunctionDef node for test_2 is popped from nodes and all child nodes are added to nodes recursively.
3. During the next iterations those are also popped one by one from nodes.
4. Without _SCOPE_END_MARKER we would reach a point at which the FunctionDef node for test_1 would be popped from nodes, not knowing that we already left the test_2 function scope -> That's why the marker is added right after the test_2 node is popped, i.e. when we reach it all child nodes have been dealt with and we should leave the current scope.

There are a few caveats, mainly that not all child nodes are actually in the child scope (e.g. default arguments which are evaluated in the parent scope), but it doesn't really matter here as those don't usually contain assert statements. So we don't need to handle them specifically.

Hope that at least somewhat makes sense.

pytest/src/_pytest/assertion/rewrite.py

Lines 722 to 743 in e5c81fa

	nodes: List[ast.AST] = [mod]
	while nodes:
	node = nodes.pop()
	for name, field in ast.iter_fields(node):
	if isinstance(field, list):
	new: List[ast.AST] = []
	for i, child in enumerate(field):
	if isinstance(child, ast.Assert):
	# Transform assert.
	new.extend(self.visit(child))
	else:
	new.append(child)
	if isinstance(child, ast.AST):
	nodes.append(child)
	setattr(node, name, new)
	elif (
	isinstance(field, ast.AST)
	# Don't recurse into expressions as they can't contain
	# asserts.
	and not isinstance(field, ast.expr)
	):
	nodes.append(field)

[nicoddemus]

Thanks!

[nicoddemus]

Also, I understand this closes #11115 as well?

[cdce8p]

Also, I understand this closes #11115 as well?

Unfortunately not, that would require a larger change to the AssertionRewriter.

[nicoddemus]

Went ahead and used an Enum as a sentinel, as I recall that's more type safe than just object (which would accept anything). @bluetech is that correct?

[bluetech]

@bluetech is that correct?

Yes, object sentinels are not great for typing.

[nicoddemus]

Is the enum solution I used there OK, or do you have any suggestions to improve it?

[cdce8p]

Went ahead and used an Enum as a sentinel, as I recall that's more type safe than just object (which would accept anything).

Yeah, although it doesn't really matter much in this case IMO. nodes is only used locally in that function and we do assert isinstance(node, ast.AST) pretty early one as well.

PEP 661 would really help in these cases but that hasn't been going anywhere unfortunately.

[nicoddemus]

Should I keep the enum change? Glad to revert if people find it is not helpful/doesn't matter.

[cdce8p]

Is the enum solution I used there OK, or do you have any suggestions to improve it?

It works but as explained above, IMO it's overkill of the situation. Would be fine with either though, leaving it as is or reverting it.

[The-Compiler]

I've done

class Sentinel: pass

SCOPE_END_MARKER = Sentinel()

before which seems a bit simpler to me, but also does have some drawbacks compared to the enum thing I suppose (such as no nice repr, or there only being one Sentinel class but possibly multiple instances of it).