fix: Bad span for missing return diagnostic with walrus loop condition

[thewill-i-am]

#1792

Fixes the bad span for the "Expected return statement" diagnostic when a while loop uses a walrus condition. Given:

@guppy
def main(b: bool) -> int:
    while (b := b):
        return 0

Before

   |     while (b := b):
   |            ^^^^^^ Expected return statement

The error highlighted only the condition b := b instead of the entire while statement, and the "Consider adding a return statement if this expression is False" help note was missing entirely.

After

   |     while (b := b):
   |     ^^^^^^^^^^^^^^^
   |         return 0
   | ^^^^^^^^^^^^^^^^ Expected return statement

Note:
   |     while (b := b):
   |            ------ Consider adding a return statement if this expression is
   |                   `False`

Root cause

ExprBuilder.visit_NamedExpr desugars a walrus (b := b) into a synthetic ast.Assign appended to the current basic block (the loop header head_bb), returning the bare Name('b') as the branch_pred. So head_bb ends up with both statements = [Assign(b := b)] and branch_pred = Name('b').

find_missing_return_point walks ancestors looking for the nearest statement-bearing block; on finding one it uses statements[-1] as the error span and only looks for a branch predicate in ancestors above that block (deliberately skipping the block itself via itertools.islice(..., 1, None)). It never noticed the statement-bearing block was itself the branch header — the predicate was right there — and returned (Assign, None), which (a) made the span cover only b := b and (b) suppressed the MissingBranch help note.

For non-walrus loops (while i < n:), head_bb.statements stays empty, so the function falls through to the second code path that returns (None, (branch_pred, 0)) — which is why those cases already worked.

Fix

Added a guard in find_missing_return_point (guppylang-internals/src/guppylang_internals/cfg/builder.py) right after computing final_statement. When the statement-bearing ancestor also owns a branch_pred and the last statement is the desugared walrus Assign whose single Name target matches the branch_pred (same id, same lineno), it returns (None, (final_statement, 0)) — routing through the same path as non-walrus loops. The caller then uses nodes[-1] (the whole while statement) as the error span and attaches MissingBranch(Assign, False) pointing at the walrus.

The guard requires five conditions simultaneously, so it only matches the walrus-in-header desugar and doesn't affect any other case — all existing no_return* golden .err snapshots are unchanged.

The function's second return-element type was widened from tuple[ast.expr, int] to tuple[ast.AST, int] since the walrus Assign (a statement, not an expression) can now flow through; MissingBranch accepts any ToSpan = ast.AST | Span.

Test

Added a regression test tests/error/misc_errors/no_return_loop_walrus.{py,err}, auto-picked-up by the parametrised tests/error/test_misc_errors.py runner. The .err golden was generated with --snapshot-update.

Verification

• uv run ruff check guppylang guppylang-internals tests — All checks passed
• uv run ruff format --check guppylang guppylang-internals tests — 315 files already formatted
• uv run mypy guppylang guppylang-internals — no issues found in 169 source files
• uv run pytest -n auto --ignore=tests/integration/test_notebooks.py — 1846 passed, 25 skipped, 8 xfailed
• The 15 tests/integration/test_notebooks.py failures are pre-existing environment issues (ModuleNotFoundError: No module named 'guppylang' in the Jupyter kernel) unrelated to this change.

Files changed

• guppylang-internals/src/guppylang_internals/cfg/builder.py — fix + type widening
• tests/error/misc_errors/no_return_loop_walrus.py — new regression test
• tests/error/misc_errors/no_return_loop_walrus.err — new golden snapshot

[croyzor]

Hi @thewill-i-am, thanks for the contribution! Sorry for the slow reply, it took me a while to get my head around this part of the codebase. This is a nice error message improvement, but I think it's not sufficient in its current form. Thinking about what other forms of surface syntax would be affected by this, both

if b:
    return

and

if True if b else False:
    return

get the improved error message, but it falls over when we have

if not b:
    return

or

if False if b else True:
    return

because the code is always suggesting to add a return if b is False, which isn't always appropriate.

I'm happy for you to consider this and submit an updated version of the PR if you're still interested 🙂

[croyzor]

I think a slightly more flexible condition here would be

isinstance(fbb_ancestor.branch_pred, ast.Name)
and isinstance(final_statement, ast.Assign)
and fbb_ancestor.branch_pred.id
in [
    tgt.id
    for tgt in final_statement.targets
    if isinstance(tgt, ast.Name)
]

because we could then catch something like

a, b = foo()
if a:
    ...

[thewill-i-am]

@croyzor

Thanks for the suggestion! I broadened the condition to support assignments with multiple or nested targets.
I used ast.walk over each assignment target and matched only ast.Name nodes with an ast.Store context. Python represents a, b = foo() as a single ast.Tuple target, so this also handles the tuple-unpacking example while avoiding false matches such as obj in obj.field = value.
I also added a regression test for:

a, b = values
if a:
    return 0

The diagnostic now highlights the if statement and points the help note at a.