await in list comprehension - #5334

Currently trying to figure out how to make extra_tests/snippets/syntax_async.py work

maybe missing implementation in this file? https://github.com/RustPython/RustPython/blob/main/vm/src/frame.rs#L977-L1047

On a sidenote, it looks like there is a current bug in branch main related to extra_tests/snippets/syntax_async.py, where if one tries to print out the list tasks or access an element of it, it panics:

thread 'main' panicked at vm/src/frame.rs:169:29:
attempt to subtract with overflow

Edit: looks like an issue already exists for this: #4859

So I fixed one issue so code like this runs:

import asyncio
async def run():
    return (None for _ in [])
asyncio.run(run(), debug=True)

However, I also need code like this to work:

import asyncio
async def run():
    x = (None for _ in [1])
    print(type(x))
    return (e for e in x)
asyncio.run(run(), debug=True)

For async list comprehensions, its ok for the generated function to be a coroutine and then just await it. But for generators its going to be more complicates since the async version, async_generator, is fundamentally different (aiter instead of iter).

It looks like in CPython, the interpreter automatically detects which a generator should be by looking at wether the comprehension element is async or not.
E.g. (await asyncio.sleep(1) for _ in [1]) is async_generator but (None for _ in [1]) is generator.

@youknowone Is there currently a way to determine if an expression needs to be compiled async (e.g. uses await somewhere) or is there a way to try compiling an expression (as sync), and if that fails, backtrack and try another way (as async)?

Checked how it compiled

>>> import ast
>>> top = ast.parse('(await asyncio.sleep(1) for _ in [1])')
>>> ast.dump(top.body[0])
"Expr(value=GeneratorExp(elt=Await(value=Call(func=Attribute(value=Name(id='asyncio', ctx=Load()), attr='sleep', ctx=Load()), args=[Constant(value=1)], keywords=[])), generators=[comprehension(target=Name(id='_', ctx=Store()), iter=List(elts=[Constant(value=1)], ctx=Load()), ifs=[], is_async=0)]))"

Expr::Await exists. But it flags nothing before calling self.compile_expression(value)?;. How about flagging it just like how compile_function_def is doing?

I have some flagging happening at the top of compile_comprehension.

        let prev_ctx = self.ctx;
        let is_async_gen = generators.iter().any(|g| g.is_async);
        let is_async = is_async_gen || prev_ctx.func == FunctionContext::AsyncFunction;

        self.ctx = CompileContext {
            loop_data: None,
            in_class: prev_ctx.in_class,
            func: if is_async {
                FunctionContext::AsyncFunction
            } else {
                FunctionContext::Function
            },
        };

        // We must have at least one generator:
        assert!(!generators.is_empty());

        let flags = bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED;
        let flags = if is_async {
            flags | bytecode::CodeFlags::IS_COROUTINE
        } else {
            flags
        };
        ```


But I think my problem right now is that the comprehension element expression can determine whether it is async or not. 

```python
async def run():
    x = (time.sleep(0.1) for _ in [1])
    y = (await asyncio.sleep(0.1) for _ in [1])
    print(type(x))
    print(type(y))

Running this with CPython3.12 returns:

<class 'generator'>
<class 'async_generator'>

In compile_function_def, it takes an is_async argument and then determines how to evaluate the body of the function. But in the case of generators, it is the generator itself and its comprehension element expression that determines if it should be an async generator or a regular one.

Accidentally pressed the wrong button - thus the close and reopen

How about this kind of approach?

diff --git a/compiler/codegen/src/compile.rs b/compiler/codegen/src/compile.rs
index 993de0f82..9f59cbd4a 100644
--- a/compiler/codegen/src/compile.rs
+++ b/compiler/codegen/src/compile.rs
@@ -883,7 +883,7 @@ impl Compiler {
                                 statement.location(),
                             ));
                         }
-                        self.compile_expression(v)?;
+                        self.compile_expression(v, self.ctx.func == FunctionContext::AsyncFunction)?;
                         self.emit_return_value();
                     }
                     None => {
@@ -2216,7 +2216,11 @@ impl Compiler {
         Ok(())
     }
 
-    fn compile_expression(&mut self, expression: &located_ast::Expr) -> CompileResult<()> {
+    fn compile_sync_expression(&mut self, expression: &located_ast::Expr) -> CompileResult<()> {
+        self.compile_expression(expression, false)
+    }
+
+    fn compile_expression(&mut self, expression: &located_ast::Expr, is_await: bool) -> CompileResult<()> {
         use located_ast::*;
         trace!("Compiling {:?}", expression);
         let location = expression.location();
@@ -2332,7 +2336,7 @@ impl Compiler {
                 if self.ctx.func != FunctionContext::AsyncFunction {
                     return Err(self.error(CodegenErrorType::InvalidAwait));
                 }
-                self.compile_expression(value)?;
+                self.compile_expression(value, true)?;
                 emit!(self, Instruction::GetAwaitable);
                 self.emit_load_const(ConstantData::None);
                 emit!(self, Instruction::YieldFrom);

I think I finally got all the edge cases - I ended up writing a method to determine whether an expression contains await as CPython acts differently if it does or does not contain an await, so I replicated that here.

I wrote a more comprehensive set of tests for various possible async comprehensions and tweaked & simplified the logic into two cases to pass these tests: async generators and async list/set/dict.

However, it looks like now this test is panicking:

Likely my changes have removed the TypeError, but some other part is missing to make this test work as an async generator that never awaits anything runs fine, but one that does await panics when running the frame.

    # TODO: RUSTPYTHON; TypeError: object async_generator can't be used in 'await' expression
    @unittest.expectedFailure
    def test_async_gen_expression_01(self):
        async def arange(n):
            for i in range(n):
                await asyncio.sleep(0.01)
                yield i

        def make_arange(n):
            # This syntax is legal starting with Python 3.7
            return (i * 2 async for i in arange(n))

        async def run():
            return [i async for i in make_arange(10)]

        res = self.loop.run_until_complete(run())
        self.assertEqual(res, [i * 2 for i in range(10)])

I'll continue looking into it.

I only renamed the test file to be aligned with syntax_async.py