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

[Ontonator]

Describe the Bug

import asyncio
import typing

async def some_async_func(x: int) -> bool:
    return x % 2 == 0

async def main() -> None:
    generator = (x for x in [1, 2, 3] if await some_async_func(x) == True)
    typing.reveal_type(generator)
    assert await anext(generator) == 2

# asyncio nonsense to make this work in the sandbox
asyncio.ensure_future(main(), loop=asyncio.get_event_loop())

In the code snippet above, Pyrefly infers the type of generator to be Generator[int, None, None] rather than AsyncGenerator[int, None].

I previously raised #1299 (very similar, but note the change to the condition expression in the generator) for the case where the await expression comprised the full condition of the generator, and that has since been fixed. However, I have now noticed that it doesn't work in the case where the await is nested inside the condition expression. I skimmed the code to try to understand what was going on, and I suspect it is because the visitor that visits that expression:

pyrefly/crates/pyrefly_python/src/ast.rs

Lines 331 to 335 in cd06d67

	expr.visit(&mut |node: &Expr| {
	if matches!(node, Expr::Await(_)) {
	found = true;
	}
	});

doesn't recurse like the example does in the documentation of visit:

pyrefly/crates/pyrefly_util/src/visit.rs

Lines 47 to 52 in d6f28be

	/// fn f(x: &Expr, res: &mut Vec<i32>) {
	/// if let Expr::Constant(x) = x {
	/// res.push(*x)
	/// }
	/// x.recurse(|x| f(x, res));
	/// }

Sandbox Link

https://pyrefly.org/sandbox/?project=N4IgZglgNgpgziAXKOBDAdgEwEYHsAeAdAA4CeSIEAtsbgE4AuABKnKegMYS4A661tRkwaliEdAHM%2BfVuw5NMMMEzi4qMAPqzOGsAFdOACnyIm4hgEomAWgB8TPLiiI%2BTN0zowGeuuib4mAFImACYmAF5wpgAGaXRteUVlKlRxQys7JgA5XHQYFz93CRg8ulQGegimYyYwSoDxJgBtAEYAGlCOgGYAXTNlVAB3VOZVdS02HX0jfCtIpgAVOj0YC1d3ETFJQk8ANxhUKA1NmENi0vL6NcK3VjgYISGRljz8BjOSh8u6OaiQuIAxCxJlxcEx0Ll7uh7sIwSkANYwYQACwgcCYg3o8LMfgYyKRaCweHwMhB3EIJTgPk0%2Bm8nkMKTSFg6UFwuGI4QS5OKDA0MH26F5rPZ6QsIDaIDInjAUFIhAqVCgFCBAAVSNLZSoMDgCEwOLlIBIfOVuOhCHwgQBlGBI5EMBjEOCIAD0zqlSllhHoEmdJWdmFwHDgzv16ENxoYpudtUqqF2qSgqGwsD1BogRrKkdyTHZWeh5vQZDxuWs%2BzocFNVR4IC6hBa-xAfBAAF8JagOJH9gAxaAwCiEnVEMgtoA

(Only applicable for extension issues) IDE Information

No response

[yangdanny97]

Thanks for reporting!

[AryanBagade]

@yangdanny97 can i work on this???

[AryanBagade]

will submit a pr soon?

[yangdanny97]

assigned, thanks