Propagate walrus narrowing from nested expressions

[Endika]

Fixes #19430.

#19038 added narrowing for walrus assignments at the top level of a condition,
and for the (foo := val).attr / (foo := val)[key] forms reached through
refine_parent_types. Its helper _propagate_walrus_assignments notes that it
"only considers nested assignment exprs, does not recurse into other types" and
that a dedicated visitor could be added later. This is that visitor.

A walrus nested anywhere else in the condition was not propagated:

def main(cond: bool, n: int) -> None:
    woo = None
    if cond and (woo := 5) + n:
        reveal_type(woo)  # was int | None, now int

How it works

This applies only to the right operand of and and or. Everywhere else the
operand is always evaluated, so the binder already carries the assignment, and
adding a map entry there would only widen the result — see the mypy_primer note
below.

WalrusAssignmentCollector walks that operand and collects the walrus
assignments it is guaranteed to perform. Traversal stops wherever evaluation is
not guaranteed, or happens in another scope:

• and / or, which short-circuit
• the branches of a conditional expression (its condition is still walked)
• comprehensions and generator expressions
• lambda bodies
• the third and later operands of a chained comparison, since a < b < c stops
  at the first false comparison

Every collected assignment has therefore taken place once the condition has
been evaluated, whatever value it produced, so the assigned type is added to
both the if and the else map.

Deciding which branches the narrowing actually reaches is then left to the
existing map algebra, which already models exactly this. For
cond and (woo := 5) + n:

• and_conditional_maps(left_if, right_if) keeps the entry, because it is in
  the right operand's if map — narrowed in the if branch.
• or_conditional_maps(left_else, right_else, coalesce_any=True) drops it,
  because it is not in the left operand's else map — not narrowed in the else
  branch.

and symmetrically for or, which narrows in the else branch instead. This
matches the framing in the issue: the narrowing applies where the assignment is
reached on every path, i.e. for and but not for or.

Nested short-circuiting operators are not walked into, but they are still
handled, because find_isinstance_check recurses into them itself. There is a
test for that.

An existing entry for the same target is never overwritten, so narrowing from
the condition itself — truthiness, isinstance, is not None — is not weakened
to the bare assigned type. That is the main risk in this change, and it has its
own test case.

Tests

Three cases in check-inference.test, next to the ones from #19038:

• testInferWalrusAssignmentNestedInCondition — binary operation, call,
  isinstance, comparison, unary operation, nested walrus, or narrowing the
  else branch, and a nested and.
• testInferWalrusAssignmentNestedInConditionNotAlwaysEvaluated — the branches
  of a conditional expression, a comprehension and a chained comparison operand
  are not narrowed, while a walrus in a conditional expression's condition is.
• testInferWalrusAssignmentDoesNotWeakenNarrowing — truthiness, isinstance
  and is not None narrowing survive.

Only the first of those fails on master. The other two are guards against this
change over-narrowing rather than regression tests for the issue, so I checked
what they actually catch by breaking each restriction in turn and seeing whether
a test noticed:

Restriction removed	Caught by
Propagation limited to and / or right operands	DoesNotWeakenNarrowing
Existing narrowing not overwritten	DoesNotWeakenNarrowing
Conditional expression branches skipped	NotAlwaysEvaluated
Comprehensions skipped	NotAlwaysEvaluated
Chained comparison operands after the second skipped	NotAlwaysEvaluated
Lambda bodies skipped	nothing

Every condition in those two cases puts the walrus on the right of an and, so
that the code path is actually reached.

The lambda restriction has no test because it is unobservable: a walrus in a
lambda body binds in the lambda's scope, so the outer variable is untouched
either way, and the whole suite passes with the restriction removed. I kept it
rather than dropping it, since walking into a body that is not evaluated would be
wrong on its own terms.

Full suites run locally on Python 3.14 / Linux, uncompiled: testcheck 8203
passed, and testfinegrained / testmerge / testtransform / testdeps /
testpythoneval 1453 passed. Self-check clean.

A pre-existing soundness bug this turned up

While testing the chained-comparison restriction I found that outside an and,
a walrus in a lazily evaluated chain operand is narrowed as though it had always
run:

class C:
    def __lt__(self, other: "C") -> bool: ...

def chain(a: C, b: C, v: C) -> None:
    woo = None
    if a < b < (woo := v):
        reveal_type(woo)  # C -- correct, a < b was true so the walrus ran
    else:
        reveal_type(woo)  # C -- unsound, a < b may have been false

If a < b is false the chain short-circuits, the walrus never runs, and woo is
still None in the else branch. Writing the same thing with and gives the
correct C | None, and putting the walrus in the first pair (a < (woo := v) < c)
correctly gives C.

The cause is not narrowing. visit_comparison_expr documents that a < b > c is
"check as a < b and b > c", but it accepts the operands in a plain loop with no
binder frame per pair, whereas check_boolean_op frames each side. So the
assignment is recorded unconditionally. Nothing narrows C here, which is how I
ruled out comparison_type_narrowing_helper.

I confirmed against master that this predates this PR and is not affected by it.
It is also why the restriction in the collector has to be tested through an
and, where the binder does not mask it. check_chained_comparison in
testInferWalrusAssignmentNestedInConditionNotAlwaysEvaluated pins the current
behaviour with a comment, so a fix will show up as a test change rather than
silently.

I have deliberately not fixed it here. It is the opposite defect — narrowing
applied where it should not be — it lives in expression checking rather than in
narrowing, and giving chained comparisons per-pair binder frames is a change with
real blast radius that deserves reviewing on its own. Happy to open an issue or a
PR for it, whichever you prefer.

The mypy_primer diff

The first version of this hooked the propagation into find_isinstance_check
itself, and mypy_primer caught one diff for it, a new error in rotki. It was a
genuine regression rather than a true positive, and it is worth spelling out
because it shaped the design:

def d(m: dict[str, tuple[str, int]], jgi: dict[str, str]) -> None:
    mgid = jgi.get("x")            # declared str | None
    if (mgid := m["y"][0]) == "z":
        pass
    reveal_type(mgid)              # master: str, first version: str | None

Inside both branches the type was str either way; the widening happened when
the branches were joined. With no map entry the binder simply keeps the type from
the assignment. With an entry, the join goes through join_simple bounded by the
target's declaration, which here is wider than what was assigned.

The operand in that condition is always evaluated, so the entry was not buying
anything: if (woo := 5) + n: already narrows correctly on master without this
PR. Restricting the propagation to the right operand of and and or removes
the regression and keeps every case the PR is for.
check_declaration_wider_than_assignment covers it.

Performance

The collector only runs for the right operand of and and or, and returns
immediately when nothing was collected.

Self-check with a cold cache, three runs each, uncompiled:

	run 1	run 2	run 3
this branch	14.50s	14.31s	14.35s
master	14.42s	14.51s	14.18s

No measurable difference on the median. I have not measured a compiled build, so
please treat these as indicative.

Known limitations

• Nested short-circuiting operators reached through another expression, such as
  truthy(a and (woo := 5)), are not narrowed. Handling them would mean
  deciding which branch each side belongs to inside the collector, which is
  what find_isinstance_check already does; it seemed better not to duplicate
  that.
• The in narrowing from the original report in Walrus narrowing not propagated from inside expression #19430 is Type not narrowed correctly for if X in collection #3229 and is
  unaffected by this.

LLM disclosure

Per the contributing guidelines: this text was written with LLM assistance. I am a
first-time contributor here, so I understand that may weigh against the PR —
please say so if you would rather not take it on that basis and I will close it
without any fuss. I have read the surrounding code and can explain and defend
every decision above. Three of them came out of verification rather than from
the first draft, and are worth stating explicitly:

• An early version guarded the type maps with if type_map is None. Self-check
  flagged it as unreachable, correctly: TypeMap is dict[Expression, Type]
  with no Optional, and conditional_types_to_typemaps returns {} rather
  than None. The guard was dead code and is gone. Incidentally the docstring
  of find_isinstance_check still claims it "can return None, None", which
  looks stale — I left it alone as it is out of scope here.
• I expected a < b < (woo := 5) not to narrow in the if branch, and wrote a
  test asserting that. It was wrong: entering the if branch implies a < b was
  true, so the third operand was evaluated after all. Chasing that down is what
  surfaced the soundness bug above.
• I expected cond and (1 and (woo := 5)) not to narrow, since the collector
  does not walk into nested and. It does narrow, because
  find_isinstance_check recurses into it. The test now records the real
  behaviour.
• The mypy_primer diff on the first push was a real regression in this PR, not a
  true positive. Finding that out is what led to restricting the propagation to
  short-circuiting operands, which is a better design than what I started with.
• My first diagnosis of the chained comparison bug blamed collapse_walrus in
  comparison_type_narrowing_helper. That was wrong, and the second commit here
  corrects the comment: reproducing it with a class whose __lt__ returns plain
  bool shows nothing narrows the operand at all, which ruled that path out and
  pointed at the missing binder frame instead.

[github-actions]

According to mypy_primer, this change doesn't affect type check results on a corpus of open source code. ✅

[Endika]

the primer diff was a real regression in this PR rather than a true positive. Fixed in fa6ef09, and the PR description is updated.

Reduced, with rotki's declarations kept:

def d(m: dict[str, tuple[str, int]], jgi: dict[str, str]) -> None:
    mgid = jgi.get("x")            # declared str | None
    if (mgid := m["y"][0]) == "z":
        pass
    reveal_type(mgid)              # master: str, before the fix: str | None

Inside both branches the revealed type was str either way — the widening
happened at the join. With no map entry the binder just keeps the type from the
assignment; with an entry the join goes through join_simple bounded by the
target's declaration, which is wider here than what was assigned.

The operand in that condition is always evaluated, so the entry was not buying
anything in the first place: if (woo := 5) + n: already narrows correctly on
master. So the propagation now applies only to the right operand of and and
or, which is exactly where the binder cannot carry the assignment. Every case
the PR is for still works, and check_declaration_wider_than_assignment covers
the regression.

Two side effects of the fix worth mentioning:

• The tests for the restrictions in the collector had to move behind an and,
  since that is now the only path that reaches it. That made them stricter: the
  chained-comparison restriction previously had no test that could fail, and now
  it does.
• The collector runs on far fewer expressions, so the earlier ~1% self-check
  cost is gone — three cold runs each give a median of 14.35s on this branch
  against 14.42s on master, uncompiled.