Skip to content

Commit 09f335b

Browse files
mikasoukhovmikasoukhov
committed
Implement PEP 492: async/await support
- Tokenizer/Parser: async def, async for, async with, await keywords - AST nodes: AwaitExpression, AsyncForStatement, AsyncWithStatement - Runtime: PythonCoroutine, CoroutineWrapper types - Code generation: coroutines reuse generator state machine via yield from desugaring (await → yield from expr.__await__()) - Fix GeneratorRewriter VisitExtension to reduce one level at a time, preventing "must be reducible node" with DebugInfoRemovalExpression Verified against CPython 3.14: 20/20 comparison tests identical.
1 parent 29308e7 commit 09f335b

16 files changed

Lines changed: 803 additions & 45 deletions

‎src/core/IronPython/Compiler/Ast/AstMethods.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ internal static class AstMethods {
7979
public static readonly MethodInfo PushFrame = GetMethod((Func<CodeContext, FunctionCode, List<FunctionStack>>)PythonOps.PushFrame);
8080
public static readonly MethodInfo FormatString = GetMethod((Func<CodeContext, string, object, string>)PythonOps.FormatString);
8181
public static readonly MethodInfo GeneratorCheckThrowableAndReturnSendValue = GetMethod((Func<object, object>)PythonOps.GeneratorCheckThrowableAndReturnSendValue);
82+
public static readonly MethodInfo MakeCoroutine = GetMethod((Func<PythonFunction, MutableTuple, object, PythonCoroutine>)PythonOps.MakeCoroutine);
8283

8384
// builtins
8485
public static readonly MethodInfo Format = GetMethod((Func<CodeContext, object, string, string>)PythonOps.Format);

‎src/core/IronPython/Compiler/Ast/AsyncForStatement.cs‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the Apache 2.0 License.
3+
// See the LICENSE file in the project root for more information.
4+
5+
#nullable enable
6+
7+
using System.Threading;
8+
9+
using Microsoft.Scripting;
10+
using MSAst = System.Linq.Expressions;
11+
12+
namespace IronPython.Compiler.Ast {
13+
14+
/// <summary>
15+
/// Represents an async for statement.
16+
/// Desugared to Python AST that uses __aiter__ and await __anext__().
17+
/// </summary>
18+
public class AsyncForStatement : Statement, ILoopStatement {
19+
private static int _counter;
20+
private Statement? _desugared;
21+
22+
public AsyncForStatement(Expression left, Expression list, Statement body, Statement? @else) {
23+
Left = left;
24+
List = list;
25+
Body = body;
26+
Else = @else;
27+
}
28+
29+
public int HeaderIndex { private get; set; }
30+
31+
public Expression Left { get; }
32+
33+
public Expression List { get; set; }
34+
35+
public Statement Body { get; set; }
36+
37+
public Statement? Else { get; }
38+
39+
MSAst.LabelTarget ILoopStatement.BreakLabel { get; set; } = null!;
40+
41+
MSAst.LabelTarget ILoopStatement.ContinueLabel { get; set; } = null!;
42+
43+
/// <summary>
44+
/// Build the desugared tree. Called during Walk when Parent and IndexSpan are available.
45+
/// </summary>
46+
private Statement BuildDesugared() {
47+
var parent = Parent;
48+
var span = IndexSpan;
49+
var id = Interlocked.Increment(ref _counter);
50+
51+
// async for TARGET in ITER:
52+
// BLOCK
53+
// else:
54+
// ELSE_BLOCK
55+
//
56+
// desugars to:
57+
//
58+
// __aiter = ITER.__aiter__()
59+
// __running = True
60+
// while __running:
61+
// try:
62+
// TARGET = await __aiter.__anext__()
63+
// except StopAsyncIteration:
64+
// __running = False
65+
// else:
66+
// BLOCK
67+
// else:
68+
// ELSE_BLOCK
69+
70+
var iterName = $"__asyncfor_iter{id}";
71+
var runningName = $"__asyncfor_running{id}";
72+
73+
// Helper to create nodes with proper parent and span
74+
NameExpression MakeName(string name) {
75+
var n = new NameExpression(name) { Parent = parent };
76+
n.IndexSpan = span;
77+
return n;
78+
}
79+
80+
T WithSpan<T>(T node) where T : Node {
81+
node.IndexSpan = span;
82+
return node;
83+
}
84+
85+
// _iter = ITER.__aiter__()
86+
var aiterAttr = WithSpan(new MemberExpression(List, "__aiter__") { Parent = parent });
87+
var aiterCall = WithSpan(new CallExpression(aiterAttr, null, null) { Parent = parent });
88+
var assignIter = WithSpan(new AssignmentStatement(new Expression[] { MakeName(iterName) }, aiterCall) { Parent = parent });
89+
90+
// running = True
91+
var trueConst = new ConstantExpression(true) { Parent = parent }; trueConst.IndexSpan = span;
92+
var assignRunning = WithSpan(new AssignmentStatement(new Expression[] { MakeName(runningName) }, trueConst) { Parent = parent });
93+
94+
// TARGET = await __aiter.__anext__()
95+
var anextAttr = WithSpan(new MemberExpression(MakeName(iterName), "__anext__") { Parent = parent });
96+
var anextCall = WithSpan(new CallExpression(anextAttr, null, null) { Parent = parent });
97+
var awaitNext = new AwaitExpression(anextCall);
98+
var assignTarget = WithSpan(new AssignmentStatement(new Expression[] { Left }, awaitNext) { Parent = parent });
99+
100+
// except StopAsyncIteration: __running = False
101+
var falseConst = new ConstantExpression(false) { Parent = parent }; falseConst.IndexSpan = span;
102+
var stopRunning = WithSpan(new AssignmentStatement(
103+
new Expression[] { MakeName(runningName) }, falseConst) { Parent = parent });
104+
var handler = WithSpan(new TryStatementHandler(
105+
MakeName("StopAsyncIteration"),
106+
null!,
107+
WithSpan(new SuiteStatement(new Statement[] { stopRunning }) { Parent = parent })
108+
) { Parent = parent });
109+
handler.HeaderIndex = span.End;
110+
111+
// try/except/else block
112+
var tryExcept = WithSpan(new TryStatement(
113+
assignTarget,
114+
new[] { handler },
115+
WithSpan(new SuiteStatement(new Statement[] { Body }) { Parent = parent }),
116+
null!
117+
) { Parent = parent });
118+
tryExcept.HeaderIndex = span.End;
119+
120+
// while __running: try/except/else
121+
var whileStmt = new WhileStatement(MakeName(runningName), tryExcept, Else);
122+
whileStmt.SetLoc(GlobalParent, span.Start, span.End, span.End);
123+
whileStmt.Parent = parent;
124+
125+
var suite = WithSpan(new SuiteStatement(new Statement[] { assignIter, assignRunning, whileStmt }) { Parent = parent });
126+
return suite;
127+
}
128+
129+
public override MSAst.Expression Reduce() {
130+
return _desugared!.Reduce();
131+
}
132+
133+
public override void Walk(PythonWalker walker) {
134+
if (walker.Walk(this)) {
135+
// Build the desugared tree on first walk (when Parent and IndexSpan are set)
136+
if (_desugared == null) {
137+
_desugared = BuildDesugared();
138+
}
139+
_desugared.Walk(walker);
140+
}
141+
walker.PostWalk(this);
142+
}
143+
144+
internal override bool CanThrow => true;
145+
}
146+
}

‎src/core/IronPython/Compiler/Ast/AsyncWithStatement.cs‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the Apache 2.0 License.
3+
// See the LICENSE file in the project root for more information.
4+
5+
#nullable enable
6+
7+
using Microsoft.Scripting;
8+
using MSAst = System.Linq.Expressions;
9+
10+
using AstUtils = Microsoft.Scripting.Ast.Utils;
11+
12+
namespace IronPython.Compiler.Ast {
13+
using Ast = MSAst.Expression;
14+
15+
/// <summary>
16+
/// Represents an async with statement.
17+
/// Desugared to Python AST that uses await on __aenter__ and __aexit__.
18+
/// </summary>
19+
public class AsyncWithStatement : Statement {
20+
private Statement? _desugared;
21+
22+
public AsyncWithStatement(Expression contextManager, Expression? var, Statement body) {
23+
ContextManager = contextManager;
24+
Variable = var;
25+
Body = body;
26+
}
27+
28+
public int HeaderIndex { private get; set; }
29+
30+
public Expression ContextManager { get; }
31+
32+
public new Expression? Variable { get; }
33+
34+
public Statement Body { get; }
35+
36+
/// <summary>
37+
/// Build the desugared tree. Called during Walk when Parent and IndexSpan are available.
38+
/// </summary>
39+
private Statement BuildDesugared() {
40+
var parent = Parent;
41+
var span = IndexSpan;
42+
43+
// async with EXPR as VAR:
44+
// BLOCK
45+
//
46+
// desugars to:
47+
//
48+
// mgr = EXPR
49+
// try:
50+
// VAR = await mgr.__aenter__() (or just await mgr.__aenter__())
51+
// BLOCK
52+
// finally:
53+
// await mgr.__aexit__(None, None, None)
54+
55+
// Helper to create nodes with proper parent and span
56+
NameExpression MakeName(string name) {
57+
var n = new NameExpression(name) { Parent = parent };
58+
n.IndexSpan = span;
59+
return n;
60+
}
61+
62+
// mgr = EXPR
63+
var assignMgr = new AssignmentStatement(new Expression[] { MakeName("__asyncwith_mgr") }, ContextManager) { Parent = parent };
64+
assignMgr.IndexSpan = span;
65+
66+
// await mgr.__aenter__()
67+
var aenterAttr = new MemberExpression(MakeName("__asyncwith_mgr"), "__aenter__") { Parent = parent };
68+
aenterAttr.IndexSpan = span;
69+
var aenterCall = new CallExpression(aenterAttr, null, null) { Parent = parent };
70+
aenterCall.IndexSpan = span;
71+
var awaitEnter = new AwaitExpression(aenterCall);
72+
73+
Statement bodyStmt;
74+
if (Variable != null) {
75+
// VAR = await value; BLOCK
76+
var assignVar = new AssignmentStatement(new Expression[] { Variable }, awaitEnter) { Parent = parent };
77+
assignVar.IndexSpan = span;
78+
bodyStmt = new SuiteStatement(new Statement[] { assignVar, Body }) { Parent = parent };
79+
} else {
80+
var exprStmt = new ExpressionStatement(awaitEnter) { Parent = parent };
81+
exprStmt.IndexSpan = span;
82+
bodyStmt = new SuiteStatement(new Statement[] { exprStmt, Body }) { Parent = parent };
83+
}
84+
85+
// await mgr.__aexit__(None, None, None)
86+
var aexitAttr = new MemberExpression(MakeName("__asyncwith_mgr"), "__aexit__") { Parent = parent };
87+
aexitAttr.IndexSpan = span;
88+
var none1 = new ConstantExpression(null) { Parent = parent }; none1.IndexSpan = span;
89+
var none2 = new ConstantExpression(null) { Parent = parent }; none2.IndexSpan = span;
90+
var none3 = new ConstantExpression(null) { Parent = parent }; none3.IndexSpan = span;
91+
var aexitCallNormal = new CallExpression(aexitAttr,
92+
new Expression[] { none1, none2, none3 }, null) { Parent = parent };
93+
aexitCallNormal.IndexSpan = span;
94+
var awaitExitNormal = new AwaitExpression(aexitCallNormal);
95+
96+
// try/finally: await __aexit__ on normal exit
97+
var finallyExprStmt = new ExpressionStatement(awaitExitNormal) { Parent = parent };
98+
finallyExprStmt.IndexSpan = span;
99+
var tryFinally = new TryStatement(bodyStmt, null, null, finallyExprStmt) { Parent = parent };
100+
tryFinally.IndexSpan = span;
101+
tryFinally.HeaderIndex = span.End;
102+
103+
var suite = new SuiteStatement(new Statement[] { assignMgr, tryFinally }) { Parent = parent };
104+
suite.IndexSpan = span;
105+
return suite;
106+
}
107+
108+
public override MSAst.Expression Reduce() {
109+
return _desugared!.Reduce();
110+
}
111+
112+
public override void Walk(PythonWalker walker) {
113+
if (walker.Walk(this)) {
114+
// Build the desugared tree on first walk (when Parent and IndexSpan are set)
115+
if (_desugared == null) {
116+
_desugared = BuildDesugared();
117+
}
118+
_desugared.Walk(walker);
119+
}
120+
walker.PostWalk(this);
121+
}
122+
}
123+
}

‎src/core/IronPython/Compiler/Ast/AwaitExpression.cs‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the Apache 2.0 License.
3+
// See the LICENSE file in the project root for more information.
4+
5+
#nullable enable
6+
7+
using MSAst = System.Linq.Expressions;
8+
9+
using AstUtils = Microsoft.Scripting.Ast.Utils;
10+
11+
namespace IronPython.Compiler.Ast {
12+
using Ast = MSAst.Expression;
13+
14+
/// <summary>
15+
/// Represents an await expression. Implemented as yield from expr.__await__().
16+
/// </summary>
17+
public class AwaitExpression : Expression {
18+
private readonly Statement _statement;
19+
private readonly NameExpression _result;
20+
21+
public AwaitExpression(Expression expression) {
22+
Expression = expression;
23+
24+
// await expr is equivalent to yield from expr.__await__()
25+
// We build: __awaitprefix_EXPR = expr; yield from __awaitprefix_EXPR.__await__(); __awaitprefix_r = __yieldfromprefix_r
26+
var parent = expression.Parent;
27+
28+
var awaitableExpr = new NameExpression("__awaitprefix_EXPR") { Parent = parent };
29+
var getAwait = new MemberExpression(awaitableExpr, "__await__") { Parent = parent };
30+
var callAwait = new CallExpression(getAwait, null, null) { Parent = parent };
31+
var yieldFrom = new YieldFromExpression(callAwait);
32+
33+
Statement s1 = new AssignmentStatement(new Expression[] { new NameExpression("__awaitprefix_EXPR") { Parent = parent } }, expression) { Parent = parent };
34+
Statement s2 = new ExpressionStatement(yieldFrom) { Parent = parent };
35+
Statement s3 = new AssignmentStatement(
36+
new Expression[] { new NameExpression("__awaitprefix_r") { Parent = parent } },
37+
new NameExpression("__yieldfromprefix_r") { Parent = parent }
38+
) { Parent = parent };
39+
40+
_statement = new SuiteStatement(new Statement[] { s1, s2, s3 }) { Parent = parent };
41+
42+
_result = new NameExpression("__awaitprefix_r") { Parent = parent };
43+
}
44+
45+
public Expression Expression { get; }
46+
47+
public override MSAst.Expression Reduce() {
48+
return Ast.Block(
49+
typeof(object),
50+
_statement,
51+
AstUtils.Convert(_result, typeof(object))
52+
).Reduce();
53+
}
54+
55+
public override void Walk(PythonWalker walker) {
56+
if (walker.Walk(this)) {
57+
Expression?.Walk(walker);
58+
_statement.Walk(walker);
59+
_result.Walk(walker);
60+
}
61+
walker.PostWalk(this);
62+
}
63+
64+
public override string NodeName => "await expression";
65+
}
66+
}

‎src/core/IronPython/Compiler/Ast/FunctionDefinition.cs‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ internal override int KwOnlyArgCount {
117117

118118
public Expression ReturnAnnotation { get; internal set; }
119119

120-
internal override bool IsGeneratorMethod => IsGenerator;
120+
internal override bool IsGeneratorMethod => IsGenerator || IsAsync;
121121

122122
/// <summary>
123123
/// The function is a generator
@@ -182,10 +182,14 @@ internal override FunctionAttributes Flags {
182182
fa |= FunctionAttributes.ContainsTryFinally;
183183
}
184184

185-
if (IsGenerator) {
185+
if (IsGenerator || IsAsync) {
186186
fa |= FunctionAttributes.Generator;
187187
}
188188

189+
if (IsAsync) {
190+
fa |= FunctionAttributes.Coroutine;
191+
}
192+
189193
if (GeneratorStop) {
190194
fa |= FunctionAttributes.GeneratorStop;
191195
}
@@ -353,8 +357,8 @@ internal MSAst.Expression MakeFunctionExpression() {
353357
annotations
354358
)
355359
),
356-
IsGenerator ?
357-
(MSAst.Expression)new PythonGeneratorExpression(code, GlobalParent.PyContext.Options.CompilationThreshold) :
360+
(IsGenerator || IsAsync) ?
361+
(MSAst.Expression)new PythonGeneratorExpression(code, GlobalParent.PyContext.Options.CompilationThreshold, IsAsync) :
358362
(MSAst.Expression)code
359363
);
360364
} else {
@@ -652,10 +656,10 @@ private LightLambdaExpression CreateFunctionLambda() {
652656
new SourceSpan(new SourceLocation(0, start.Line, start.Column), new SourceLocation(0, start.Line, int.MaxValue))));
653657

654658

655-
// For generators, we need to do a check before the first statement for Generator.Throw() / Generator.Close().
659+
// For generators/coroutines, we need to do a check before the first statement for Generator.Throw() / Generator.Close().
656660
// The exception traceback needs to come from the generator's method body, and so we must do the check and throw
657661
// from inside the generator.
658-
if (IsGenerator) {
662+
if (IsGenerator || IsAsync) {
659663
MSAst.Expression s1 = YieldExpression.CreateCheckThrowExpression(SourceSpan.None);
660664
statements.Add(s1);
661665
}

0 commit comments

Comments
 (0)