Get enough working to pass jinja2 async tests

Author: Adam Hrbac

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/asyncio/AsyncGenSendBuiltins.java

@@ -1,7 +1,48 @@
+/*
+ * Copyright (c) 2023, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * The Universal Permissive License (UPL), Version 1.0
+ *
+ * Subject to the condition set forth below, permission is hereby granted to any
+ * person obtaining a copy of this software, associated documentation and/or
+ * data (collectively the "Software"), free of charge and under any and all
+ * copyright rights in the Software, and any and all patent rights owned or
+ * freely licensable by each licensor hereunder covering either (i) the
+ * unmodified Software as contributed to or provided by such licensor, or (ii)
+ * the Larger Works (as defined below), to deal in both
+ *
+ * (a) the Software, and
+ *
+ * (b) any piece of software and/or hardware listed in the lrgrwrks.txt file if
+ * one is included with the Software each a "Larger Work" to which the Software
+ * is contributed by such licensors),
+ *
+ * without restriction, including without limitation the rights to copy, create
+ * derivative works of, display, perform, and distribute the Software and make,
+ * use, sell, offer for sale, import, export, have made, and have sold the
+ * Software and the Larger Work(s), and to sublicense the foregoing rights on
+ * either these or other terms.
+ *
+ * This license is subject to the following condition:
+ *
+ * The above copyright notice and either this complete permission notice or at a
+ * minimum a reference to the UPL must be included in all copies or substantial
+ * portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
 package com.oracle.graal.python.builtins.objects.asyncio;
 
 import static com.oracle.graal.python.builtins.objects.asyncio.PAsyncGenASend.AwaitableState;
 import static com.oracle.graal.python.nodes.SpecialMethodNames.J___AWAIT__;
+import static com.oracle.graal.python.nodes.SpecialMethodNames.J___ITER__;
 import static com.oracle.graal.python.nodes.SpecialMethodNames.J___NEXT__;
 
 import java.util.List;
@@ -12,8 +53,11 @@
 import com.oracle.graal.python.builtins.PythonBuiltins;
 import com.oracle.graal.python.builtins.objects.PNone;
 import com.oracle.graal.python.builtins.objects.generator.CommonGeneratorBuiltins;
+import com.oracle.graal.python.builtins.objects.generator.PGenerator;
+import com.oracle.graal.python.nodes.ErrorMessages;
 import com.oracle.graal.python.nodes.PRaiseNode;
 import com.oracle.graal.python.nodes.function.PythonBuiltinBaseNode;
+import com.oracle.graal.python.nodes.function.PythonBuiltinNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonBinaryBuiltinNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonUnaryBuiltinNode;
 import com.oracle.graal.python.nodes.object.BuiltinClassProfiles;
@@ -64,17 +108,15 @@ public Object send(VirtualFrame frame, PAsyncGenASend self, Object sent,
                         @Cached BuiltinClassProfiles.IsBuiltinObjectProfile isStopIteration,
                         @Cached BuiltinClassProfiles.IsBuiltinObjectProfile isGenExit,
                         @Cached BuiltinClassProfiles.IsBuiltinObjectProfile isAsyncGenWrappedValue,
-                        @Cached PRaiseNode raiseStopAsyncIteration) {
+                        @Cached PRaiseNode raiseStopIteration) {
             Object result;
             if (self.getState() == AwaitableState.CLOSED) {
-                throw raiseReuse.raise(PythonBuiltinClassType.RuntimeError); // todo error msg
+                throw raiseReuse.raise(PythonBuiltinClassType.RuntimeError, ErrorMessages.CANNOT_REUSE_ASEND);
             }
 
             if (self.getState() == AwaitableState.INIT) {
                 if (self.receiver.isRunningAsync()) {
-                    throw raiseAlreadyRunning.raise(PythonBuiltinClassType.RuntimeError); // todo
-                                                                                          // error
-                                                                                          // msg
+                    throw raiseAlreadyRunning.raise(PythonBuiltinClassType.RuntimeError, ErrorMessages.AGEN_ALREADY_RUNNING);
                 }
                 if (sent == null || sent == PNone.NONE) {
                     sent = self.message;
@@ -84,27 +126,94 @@ public Object send(VirtualFrame frame, PAsyncGenASend self, Object sent,
             self.receiver.setRunningAsync(true);
             try {
                 result = send.execute(frame, self.receiver, sent);
-                if (result == null) {
-                    self.receiver.markAsFinished();
-                    self.receiver.setRunningAsync(false);
-                    self.setState(AwaitableState.CLOSED);
-                    throw raiseStopAsyncIteration.raise(PythonBuiltinClassType.StopAsyncIteration);
-                }
-                if (isAsyncGenWrappedValue.profileObject(inliningTarget, result, PythonBuiltinClassType.PAsyncGenAWrappedValue)) {
-                    self.receiver.setRunningAsync(false);
-                    self.setState(AwaitableState.CLOSED);
-                    throw raiseStopAsyncIteration.raise(PythonBuiltinClassType.StopIteration, new Object[]{((PAsyncGenWrappedValue) result).getWrapped()});
-                }
             } catch (PException e) {
-                if (isStopIteration.profileException(inliningTarget, e, PythonBuiltinClassType.StopAsyncIteration) ||
-                                isGenExit.profileException(inliningTarget, e, PythonBuiltinClassType.GeneratorExit)) {
-                    self.receiver.markAsFinished();
-                }
-                self.receiver.setRunningAsync(false);
+                self.setState(AwaitableState.CLOSED);
+                throw handleAGError(self.receiver, e, inliningTarget, isStopIteration, isGenExit);
+            }
+            try {
+                return unwrapAGYield(self.receiver, result, inliningTarget, isAsyncGenWrappedValue, raiseStopIteration);
+            } catch (PException e) {
                 self.setState(AwaitableState.CLOSED);
                 throw e;
             }
-            return result;
+        }
+    }
+
+    private static PException handleAGError(PGenerator self, PException exception,
+                    Node inliningTarget,
+                    BuiltinClassProfiles.IsBuiltinObjectProfile isStopAsyncIteration,
+                    BuiltinClassProfiles.IsBuiltinObjectProfile isGeneratorExit) {
+        if (isStopAsyncIteration.profileException(inliningTarget, exception, PythonBuiltinClassType.StopAsyncIteration) ||
+                        isGeneratorExit.profileException(inliningTarget, exception, PythonBuiltinClassType.GeneratorExit)) {
+            self.markAsFinished();
+        }
+        self.setRunningAsync(false);
+        return exception;
+    }
+
+    private static Object unwrapAGYield(PGenerator self, Object result,
+                    Node inliningTarget,
+                    BuiltinClassProfiles.IsBuiltinObjectProfile isAGWrappedValue,
+                    PRaiseNode raise) {
+        if (isAGWrappedValue.profileObject(inliningTarget, result, PythonBuiltinClassType.PAsyncGenAWrappedValue)) {
+            self.setRunningAsync(false);
+            Object wrapped = ((PAsyncGenWrappedValue) result).getWrapped();
+            throw raise.raise(PythonBuiltinClassType.StopIteration, new Object[]{wrapped});
+        }
+        return result;
+    }
+
+    @Builtin(name = "throw", minNumOfPositionalArgs = 2, maxNumOfPositionalArgs = 4, declaresExplicitSelf = true)
+    @GenerateNodeFactory
+    public static abstract class Throw extends PythonBuiltinNode {
+        public abstract Object execute(VirtualFrame frame, PAsyncGenASend self, Object arg1, Object arg2, Object arg3);
+
+        @Specialization
+        public Object doThrow(VirtualFrame frame, PAsyncGenASend self, Object arg1, Object arg2, Object arg3,
+                        @Bind("this") Node inliningTarget,
+                        @Cached PRaiseNode raiseReuse,
+                        @Cached CommonGeneratorBuiltins.ThrowNode throwNode,
+                        @Cached BuiltinClassProfiles.IsBuiltinObjectProfile isStopIteration,
+                        @Cached BuiltinClassProfiles.IsBuiltinObjectProfile isGeneratorExit,
+                        @Cached BuiltinClassProfiles.IsBuiltinObjectProfile isAGWrappedValue,
+                        @Cached PRaiseNode raiseStopIteration) {
+            Object result;
+
+            if (self.getState() == AwaitableState.CLOSED) {
+                throw raiseReuse.raise(PythonBuiltinClassType.RuntimeError, ErrorMessages.CANNOT_REUSE_ASEND);
+            }
+            try {
+                result = throwNode.execute(frame, self.receiver, arg1, arg2, arg3);
+            } catch (PException e) {
+                self.setState(AwaitableState.CLOSED);
+                throw handleAGError(self.receiver, e, inliningTarget, isStopIteration, isGeneratorExit);
+            }
+            try {
+                return unwrapAGYield(self.receiver, result, inliningTarget, isAGWrappedValue, raiseStopIteration);
+            } catch (PException e) {
+                self.setState(AwaitableState.CLOSED);
+                throw e;
+            }
+
+        }
+    }
+
+    @Builtin(name = J___ITER__, minNumOfPositionalArgs = 1, declaresExplicitSelf = true)
+    @GenerateNodeFactory
+    public static abstract class Iter extends PythonUnaryBuiltinNode {
+        @Specialization
+        public Object iter(PAsyncGenASend self) {
+            return self;
+        }
+    }
+
+    @Builtin(name = "close", minNumOfPositionalArgs = 1, declaresExplicitSelf = true)
+    @GenerateNodeFactory
+    public static abstract class Close extends PythonUnaryBuiltinNode {
+        @Specialization
+        public Object close(PAsyncGenASend self) {
+            self.setState(AwaitableState.CLOSED);
+            return PNone.NONE;
         }
     }
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/compiler/Compiler.java

@@ -1511,8 +1511,17 @@ private Void visitComprehension(ExprTy node, String name, ComprehensionTy[] gene
             exitScope();
             makeClosure(code, 0);
             generators[0].iter.accept(this);
-            addOp(GET_ITER);
+            if (generators[0].isAsync) {
+                addOp(GET_AITER);
+            } else {
+                addOp(GET_ITER);
+            }
             addOp(CALL_COMPREHENSION);
+            if (generators[0].isAsync) {
+                addOp(GET_AWAITABLE);
+                addOp(LOAD_NONE);
+                addYieldFrom();
+            }
             return null;
         } finally {
             setLocation(savedLocation);
@@ -1527,14 +1536,34 @@ private void visitComprehensionGenerator(ComprehensionTy[] generators, int i, Ex
         } else {
             /* Create the iterator for nested iteration */
             gen.iter.accept(this);
-            addOp(GET_ITER);
+            if (gen.isAsync) {
+                addOp(GET_AITER);
+            } else {
+                addOp(GET_ITER);
+            }
         }
         Block start = new Block();
         Block ifCleanup = new Block();
         Block anchor = new Block();
+        Block asyncForTry = gen.isAsync ? new Block() : null;
+        Block asyncForExcept = gen.isAsync ? new Block() : null;
+        Block asyncForBody = gen.isAsync ? new Block() : null;
         unit.useNextBlock(start);
-        addOp(FOR_ITER, anchor);
-        gen.target.accept(this);
+        if (gen.isAsync) {
+            addOp(DUP_TOP);
+            unit.useNextBlock(asyncForTry);
+            unit.pushBlock(new BlockInfo.AsyncForLoopExit(asyncForTry, asyncForExcept));
+            addOp(GET_ANEXT);
+            addOp(GET_AWAITABLE);
+            addOp(LOAD_NONE);
+            addYieldFrom();
+            unit.popBlock();
+            unit.useNextBlock(asyncForBody);
+            gen.target.accept(this);
+        } else {
+            addOp(FOR_ITER, anchor);
+            gen.target.accept(this);
+        }
         for (ExprTy ifExpr : gen.ifs) {
             jumpIf(ifExpr, ifCleanup, false);
         }
@@ -1566,6 +1595,10 @@ private void visitComprehensionGenerator(ComprehensionTy[] generators, int i, Ex
         }
         unit.useNextBlock(ifCleanup);
         addOp(JUMP_BACKWARD, start);
+        if (gen.isAsync) {
+            unit.useNextBlock(asyncForExcept);
+            addOp(END_ASYNC_FOR);
+        }
         unit.useNextBlock(anchor);
     }
 
@@ -2096,20 +2129,19 @@ private void visitAsyncFor(StmtTy.AsyncFor node) {
         Block body = new Block();
         Block end = new Block();
         Block except = new Block();
+        Block loopTry = new Block();
         Block orelse = node.orElse != null ? new Block() : null;
         node.iter.accept(this);
         addOp(GET_AITER);
         unit.useNextBlock(head);
         unit.pushBlock(new BlockInfo.AsyncForLoop(head, end));
-        unit.pushBlock(new BlockInfo.AsyncForLoopExit(head, except));
         addOp(DUP_TOP);
         addOp(GET_ANEXT);
+        addOp(GET_AWAITABLE);
+        unit.useNextBlock(loopTry);
+        unit.pushBlock(new BlockInfo.AsyncForLoopExit(loopTry, except));
         addOp(LOAD_NONE);
         addYieldFrom();
-        addOp(JUMP_FORWARD, body);
-        unit.useNextBlock(except);
-        addOp(END_ASYNC_FOR);
-        addOp(JUMP_FORWARD, node.orElse != null ? orelse : end);
         unit.popBlock();
         unit.useNextBlock(body);
         try {
@@ -2119,6 +2151,9 @@ private void visitAsyncFor(StmtTy.AsyncFor node) {
         } finally {
             unit.popBlock();
         }
+        addOp(JUMP_FORWARD, node.orElse != null ? orelse : end);
+        unit.useNextBlock(except);
+        addOp(END_ASYNC_FOR);
         if (node.orElse != null) {
             unit.useNextBlock(orelse);
             visitSequence(node.orElse);

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/compiler/OpCodes.java

@@ -742,9 +742,9 @@ public enum OpCodes {
      * Exception handler for async for loops. If the current exception is StopAsyncIteration, handle
      * it, otherwise, reraise.
      *
-     * Pops: exception, then the async iterator
+     * Pops: exception, then the anext coroutine, then the async iterator
      */
-    END_ASYNC_FOR(0, 2, 0),
+    END_ASYNC_FOR(0, 3, 0),
 
     // with statements
     /**

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/ErrorMessages.java

@@ -1525,4 +1525,6 @@ public abstract class ErrorMessages {
     public static final TruffleString ASYNC_FOR_NO_ANEXT_INITIAL = tsLiteral("'async for' received an object from __aiter__ that does not implement __anext__: %s");
 
     public static final TruffleString ASYNC_FOR_NO_ANEXT_ITERATION = tsLiteral("'async for' requires an iterator with __anext__ method, got %s");
+    public static final TruffleString CANNOT_REUSE_ASEND = tsLiteral("cannot reuse already awaited __anext__()/asend()");
+    public static final TruffleString AGEN_ALREADY_RUNNING = tsLiteral("anext(): asynchronous generator is already running");
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeRootNode.java

@@ -2387,7 +2387,10 @@ private int bytecodeGetANext(VirtualFrame virtualFrame, boolean useCachedNodes,
     private int bytecodeEndAsyncFor(VirtualFrame virtualFrame, boolean useCachedNodes, int stackTop, Node[] localNodes, int bci) {
         EndAsyncForNode node = insertChildNode(localNodes, bci, UNCACHED_END_ASYNC_FOR, EndAsyncForNodeGen.class, NODE_END_ASYNC_FOR, useCachedNodes);
         node.execute(virtualFrame.getObject(stackTop), frameIsVisibleToPython());
-        return stackTop - 2;
+        virtualFrame.setObject(stackTop, null); // pop the exception
+        virtualFrame.setObject(stackTop - 1, null); // the coroutine that raised the exception
+        virtualFrame.setObject(stackTop - 2, null); // the async iterator
+        return stackTop - 3;
     }
 
     @BytecodeInterpreterSwitch

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/PythonOptions.java

@@ -174,8 +174,7 @@ private PythonOptions() {
     @Option(category = OptionCategory.USER, help = "Equivalent to setting the PYTHONPATH environment variable for the standard launcher. ':'-separated list of directories prefixed to the default module search path.", usageSyntax = "<path>[:<path>]", stability = OptionStability.STABLE) //
     public static final OptionKey<TruffleString> PythonPath = new OptionKey<>(T_EMPTY_STRING, TS_OPTION_TYPE);
 
-    @EngineOption
-    @Option(category = OptionCategory.USER, help = "Equivalent to setting the PYTHONIOENCODING environment variable for the standard launcher.", usageSyntax = "<Encoding>[:<errors>]", stability = OptionStability.STABLE) //
+    @EngineOption @Option(category = OptionCategory.USER, help = "Equivalent to setting the PYTHONIOENCODING environment variable for the standard launcher.", usageSyntax = "<Encoding>[:<errors>]", stability = OptionStability.STABLE) //
     public static final OptionKey<TruffleString> StandardStreamEncoding = new OptionKey<>(T_EMPTY_STRING, TS_OPTION_TYPE);
 
     @Option(category = OptionCategory.USER, help = "Remove assert statements and any code conditional on the value of __debug__.", usageSyntax = "true|false", stability = OptionStability.STABLE) //
@@ -232,8 +231,7 @@ private PythonOptions() {
                         }
                     }));
 
-    @EngineOption
-    @Option(category = OptionCategory.USER, help = "Choose the backend for the POSIX module.", usageSyntax = "java|native|llvm") //
+    @EngineOption @Option(category = OptionCategory.USER, help = "Choose the backend for the POSIX module.", usageSyntax = "java|native|llvm") //
     public static final OptionKey<TruffleString> PosixModuleBackend = new OptionKey<>(T_JAVA, TS_OPTION_TYPE);
 
     @Option(category = OptionCategory.USER, help = "Value of the --check-hash-based-pycs command line option" +
@@ -250,12 +248,10 @@ private PythonOptions() {
     @Option(category = OptionCategory.INTERNAL, help = "Set the location of C API home. Overrides any environment variables or Java options.", usageSyntax = "<path>", stability = OptionStability.STABLE) //
     public static final OptionKey<TruffleString> CAPI = new OptionKey<>(T_EMPTY_STRING, TS_OPTION_TYPE);
 
-    @EngineOption
-    @Option(category = OptionCategory.INTERNAL, help = "Expose internal sources as normal sources, so they will show up in the debugger and stacks", usageSyntax = "true|false") //
+    @EngineOption @Option(category = OptionCategory.INTERNAL, help = "Expose internal sources as normal sources, so they will show up in the debugger and stacks", usageSyntax = "true|false") //
     public static final OptionKey<Boolean> ExposeInternalSources = new OptionKey<>(false);
 
-    @EngineOption
-    @Option(category = OptionCategory.INTERNAL, help = "Print the java stacktrace. Possible modes:" +
+    @EngineOption @Option(category = OptionCategory.INTERNAL, help = "Print the java stacktrace. Possible modes:" +
                     "    1   Print Java stacktrace for Java exceptions only." +
                     "    2   Print Java stacktrace for Python exceptions only (ATTENTION: this will have a notable performance impact)." +
                     "    3   Combines 1 and 2.", usageSyntax = "1|2|3") //