Skip to content

Conversation

@Imran-S-heikh
Copy link
Contributor

Summary

Fixes #1611.

Previously, comprehensions were incorrectly inferred as synchronous Generator instead of AsyncGenerator when the await keyword was nested inside a binary expression or comparison (e.g., if await foo() == True). The previous AST visitor did not fully recurse into the condition's children in all cases.

The Fix

I updated the contains_await logic in ast.rs to use a manual iterative Depth-First Search (DFS) using a stack.

  • Correctness: This ensures we scan the entire expression tree, identifying await nodes even when deeply nested inside other expressions.
  • Robustness: By using an iterative stack instead of recursion, we avoid potential stack overflow issues on deeply nested inputs.

Test Plan

Added a regression test test_async_generator_nested_await_in_condition in yields.rs.

Verification:

  • Verified that if await foo() == True correctly results in AsyncGenerator.
  • Ran tests locally: cargo test -p pyrefly --lib test_async_generator_nested_await_in_condition (Passed)

The previous logic failed to detect 'await' when nested inside comparisons (e.g., 'if await foo() == True').

This implements an iterative DFS traversal to correctly scan the entire expression tree, including the root and nested children.

Fixes facebook#1611
@meta-cla meta-cla bot added the cla signed label Nov 21, 2025
@stroxler stroxler self-requested a review November 21, 2025 21:48
@stroxler stroxler self-assigned this Nov 21, 2025
@stroxler
Copy link
Contributor

Importing the PR failed, it looks like there's probably a rebase conflict - would you mind rebasing on main?

@stroxler
Copy link
Contributor

Oh never mind, I think a separate PR from @AryanBagade already landed with basically the same fix.

Closing since we can't really merge these, but thanks for the PR, it's the right fix!

@stroxler stroxler closed this Nov 22, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Comprehensions are inferred as Generator not AsyncGenerator when await keyword is nested in the condition

2 participants