Skip to content

Commit 74df01f

Browse files
BCSharpBCSharp
committed
Add async generator support
1 parent 62676f3 commit 74df01f

9 files changed

Lines changed: 339 additions & 25 deletions

File tree

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ internal static class AstMethods {
8383
#if FEATURE_NET_ASYNC
8484
public static readonly MethodInfo AsTaskForAwait = GetMethod((Func<object, System.Threading.Tasks.Task<object>>)PythonOps.AsTaskForAwait);
8585
public static readonly MethodInfo MakeAsyncCoroutine = GetMethod((Func<PythonFunction, System.Threading.Tasks.Task<object>, System.Threading.CancellationTokenSource, System.Runtime.CompilerServices.StrongBox<System.Exception>, PythonCoroutine>)PythonOps.MakeAsyncCoroutine);
86+
public static readonly MethodInfo MakeAsyncGenerator = GetMethod((Func<PythonFunction, System.Collections.Generic.IAsyncEnumerable<object>, System.Runtime.CompilerServices.StrongBox<object>, System.Runtime.CompilerServices.StrongBox<System.Exception>, System.Threading.CancellationTokenSource, PythonAsyncGenerator>)PythonOps.MakeAsyncGenerator);
8687
#endif
8788

8889
// builtins

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

Lines changed: 55 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,19 @@ internal override int KwOnlyArgCount {
123123
// through the generator state machine, so IsAsync no longer implies
124124
// generator-shaped emission.
125125
internal override bool IsGeneratorMethod => IsGenerator;
126+
127+
// Async-generator (PEP 525) channels, one StrongBox per async-generator instance (per stack frame
128+
// at runtime — not static). Created when lowering the async generator; declared/assigned in the
129+
// function body, captured by the generator (the body's yields read them through Parent), and handed
130+
// to the PythonAsyncGenerator wrapper, which writes them before each resume:
131+
// AsyncSendSlot — the value of `x = yield z` (asend(v); None for __anext__/async for).
132+
// AsyncThrowSlot — an exception to rethrow at the yield resume point (athrow/aclose).
133+
private MSAst.ParameterExpression _asyncSendSlot;
134+
private MSAst.ParameterExpression _asyncThrowSlot;
135+
internal MSAst.ParameterExpression AsyncSendSlot
136+
=> _asyncSendSlot ??= MSAst.Expression.Variable(typeof(System.Runtime.CompilerServices.StrongBox<object>), "$asyncSend");
137+
internal MSAst.ParameterExpression AsyncThrowSlot
138+
=> _asyncThrowSlot ??= MSAst.Expression.Variable(typeof(System.Runtime.CompilerServices.StrongBox<Exception>), "$asyncThrow");
126139
#else
127140
internal override bool IsGeneratorMethod => IsGenerator || IsAsync;
128141
#endif
@@ -372,7 +385,9 @@ internal MSAst.Expression MakeFunctionExpression() {
372385
)
373386
),
374387
#if FEATURE_NET_ASYNC
375-
IsGenerator ?
388+
// Async generators are lowered via AsyncEnumerableExpression in the body, so they must not
389+
// be wrapped as a PythonGenerator here — only plain (non-async) generators are.
390+
(IsGenerator && !IsAsync) ?
376391
(MSAst.Expression)new PythonGeneratorExpression(code, GlobalParent.PyContext.Options.CompilationThreshold, IsAsync) :
377392
(MSAst.Expression)code
378393
#else
@@ -680,7 +695,9 @@ private LightLambdaExpression CreateFunctionLambda() {
680695
// The exception traceback needs to come from the generator's method body, and so we must do the check and throw
681696
// from inside the generator.
682697
#if FEATURE_NET_ASYNC
683-
if (IsGenerator) {
698+
// Async generators have no backing PythonGenerator (they lower to IAsyncEnumerable via
699+
// AsyncEnumerableExpression), so skip the $generator.CheckThrowable() prologue for them.
700+
if (IsGenerator && !IsAsync) {
684701
MSAst.Expression s1 = YieldExpression.CreateCheckThrowExpression(SourceSpan.None);
685702
statements.Add(s1);
686703
}
@@ -718,18 +735,42 @@ private LightLambdaExpression CreateFunctionLambda() {
718735
if (IsAsync) {
719736
var cts = MSAst.Expression.Variable(typeof(System.Threading.CancellationTokenSource), "$cts");
720737
var excBox = MSAst.Expression.Variable(typeof(System.Runtime.CompilerServices.StrongBox<Exception>), "$cancelExc");
721-
body = MSAst.Expression.Block(
722-
new[] { cts, excBox },
723-
MSAst.Expression.Assign(cts, MSAst.Expression.New(typeof(System.Threading.CancellationTokenSource))),
724-
MSAst.Expression.Assign(excBox, MSAst.Expression.New(typeof(System.Runtime.CompilerServices.StrongBox<Exception>))),
725-
Ast.Call(
726-
AstMethods.MakeAsyncCoroutine,
727-
_functionParam,
728-
AstUtils.Async(Name, body,
729-
MSAst.Expression.Property(cts, nameof(System.Threading.CancellationTokenSource.Token)),
730-
excBox),
731-
cts,
732-
excBox));
738+
var ctToken = MSAst.Expression.Property(cts, nameof(System.Threading.CancellationTokenSource.Token));
739+
if (IsGenerator) {
740+
// Async generator: the body has both `await` and `yield`. Lower it to an
741+
// IAsyncEnumerable<object?> via AsyncEnumerableExpression, sharing the generator label so
742+
// the body's yields and the rewritten awaits land in one generator, then wrap it in a
743+
// PythonAsyncGenerator. The send/throw slots are per-generator StrongBoxes captured by the
744+
// generator (the body's yields read them) AND handed to the wrapper, which writes them
745+
// before each resume — see AsyncSendSlot / AsyncThrowSlot.
746+
var sendSlot = AsyncSendSlot;
747+
var throwSlot = AsyncThrowSlot;
748+
body = MSAst.Expression.Block(
749+
new[] { cts, excBox, sendSlot, throwSlot },
750+
MSAst.Expression.Assign(cts, MSAst.Expression.New(typeof(System.Threading.CancellationTokenSource))),
751+
MSAst.Expression.Assign(excBox, MSAst.Expression.New(typeof(System.Runtime.CompilerServices.StrongBox<Exception>))),
752+
MSAst.Expression.Assign(sendSlot, MSAst.Expression.New(typeof(System.Runtime.CompilerServices.StrongBox<object>))),
753+
MSAst.Expression.Assign(throwSlot, MSAst.Expression.New(typeof(System.Runtime.CompilerServices.StrongBox<Exception>))),
754+
Ast.Call(
755+
AstMethods.MakeAsyncGenerator,
756+
_functionParam,
757+
AstUtils.AsyncEnumerable(Name, body, GeneratorLabel, ctToken, excBox),
758+
sendSlot,
759+
throwSlot,
760+
cts));
761+
} else {
762+
// Plain async def: the body returns a PythonCoroutine wrapping a Task<object?>.
763+
body = MSAst.Expression.Block(
764+
new[] { cts, excBox },
765+
MSAst.Expression.Assign(cts, MSAst.Expression.New(typeof(System.Threading.CancellationTokenSource))),
766+
MSAst.Expression.Assign(excBox, MSAst.Expression.New(typeof(System.Runtime.CompilerServices.StrongBox<Exception>))),
767+
Ast.Call(
768+
AstMethods.MakeAsyncCoroutine,
769+
_functionParam,
770+
AstUtils.Async(Name, body, ctToken, excBox),
771+
cts,
772+
excBox));
773+
}
733774
}
734775
#endif
735776

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,15 @@ public ReturnStatement(Expression expression) {
1717

1818
public override MSAst.Expression Reduce() {
1919
if (Parent.IsGeneratorMethod) {
20+
#if FEATURE_NET_ASYNC
21+
// An async generator (`async def` with `yield`) lowers through the DLR generator via
22+
// AsyncEnumerableExpression, which doesn't understand IronPython's -2 "generator-return"
23+
// marker. `return` there is always bare (`return value` is a SyntaxError in async
24+
// generators) and simply ends the async iteration — map it to a YieldBreak.
25+
if (Parent is FunctionDefinition { IsAsync: true }) {
26+
return GlobalParent.AddDebugInfo(AstUtils.YieldBreak(GeneratorLabel), Span);
27+
}
28+
#endif
2029
// Reduce to a yield return with a marker of -2, this will be interpreted as a yield break with a return value
2130
return GlobalParent.AddDebugInfo(AstUtils.YieldReturn(GeneratorLabel, TransformOrConstantNull(Expression, typeof(object)), -2), Span);
2231
}

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

Lines changed: 36 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,13 @@ namespace IronPython.Compiler.Ast {
1919
// x = yield z
2020
// The return value (x) is provided by calling Generator.Send()
2121
public class YieldExpression : Expression {
22+
#if FEATURE_NET_ASYNC
23+
private static readonly System.Reflection.MethodInfo s_captureMethod
24+
= typeof(System.Runtime.ExceptionServices.ExceptionDispatchInfo).GetMethod(nameof(System.Runtime.ExceptionServices.ExceptionDispatchInfo.Capture))!;
25+
private static readonly System.Reflection.MethodInfo s_throwMethod
26+
= typeof(System.Runtime.ExceptionServices.ExceptionDispatchInfo).GetMethod(nameof(System.Runtime.ExceptionServices.ExceptionDispatchInfo.Throw), Type.EmptyTypes)!;
27+
#endif
28+
2229
public YieldExpression(Expression? expression) {
2330
Expression = expression;
2431
}
@@ -43,16 +50,40 @@ internal static MSAst.Expression CreateCheckThrowExpression(SourceSpan span) {
4350
}
4451

4552
public override MSAst.Expression Reduce() {
53+
MSAst.Expression yieldValue = Expression == null ? AstUtils.Constant(null) : AstUtils.Convert(Expression, typeof(object));
54+
55+
#if FEATURE_NET_ASYNC
56+
// An async generator (`async def` with `yield`) is lowered via AsyncEnumerableExpression and has no
57+
// backing PythonGenerator, so there is no `$generator` to call CheckThrowable() on. Instead the
58+
// resume reads two per-generator cells that PythonAsyncGenerator writes before advancing:
59+
// AsyncThrowSlot — if set (athrow/aclose), rethrow it here (preserving stack); cleared first so a
60+
// body that catches it and yields again doesn't re-throw on the next resume.
61+
// AsyncSendSlot — the value of the yield expression: the asend(v) value, or None.
62+
if (Parent is FunctionDefinition { IsAsync: true } fd) {
63+
MSAst.ParameterExpression sendSlot = fd.AsyncSendSlot;
64+
MSAst.ParameterExpression throwSlot = fd.AsyncThrowSlot;
65+
MSAst.ParameterExpression pending = Ast.Variable(typeof(Exception), "$athrow");
66+
return Ast.Block(
67+
typeof(object),
68+
new[] { pending },
69+
AstUtils.YieldReturn(GeneratorLabel, yieldValue),
70+
Ast.Assign(pending, Ast.Field(throwSlot, nameof(System.Runtime.CompilerServices.StrongBox<Exception>.Value))),
71+
Ast.Assign(Ast.Field(throwSlot, nameof(System.Runtime.CompilerServices.StrongBox<Exception>.Value)), Ast.Constant(null, typeof(Exception))),
72+
Ast.IfThen(
73+
Ast.ReferenceNotEqual(pending, Ast.Constant(null, typeof(Exception))),
74+
Ast.Call(Ast.Call(s_captureMethod, pending), s_throwMethod)),
75+
Ast.Field(sendSlot, nameof(System.Runtime.CompilerServices.StrongBox<object>.Value))
76+
);
77+
}
78+
#endif
79+
4680
// (yield z) becomes:
4781
// .comma (1) {
4882
// .void ( .yield_statement (_expression) ),
49-
// $gen.CheckThrowable() // <-- has return result from send
83+
// $gen.CheckThrowable() // <-- has return result from send
5084
// }
5185
return Ast.Block(
52-
AstUtils.YieldReturn(
53-
GeneratorLabel,
54-
Expression == null ? AstUtils.Constant(null) : AstUtils.Convert(Expression, typeof(object))
55-
),
86+
AstUtils.YieldReturn(GeneratorLabel, yieldValue),
5687
CreateCheckThrowExpression(Span) // emits ($gen.CheckThrowable())
5788
);
5889
}

‎src/core/IronPython/Runtime/AsyncGenerator.cs‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
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+
#if FEATURE_NET_ASYNC
8+
9+
using System;
10+
using System.Collections.Generic;
11+
using System.Runtime.CompilerServices;
12+
using System.Threading;
13+
using System.Threading.Tasks;
14+
15+
using Microsoft.Scripting.Runtime;
16+
17+
using IronPython.Runtime.Exceptions;
18+
using IronPython.Runtime.Operations;
19+
using IronPython.Runtime.Types;
20+
21+
namespace IronPython.Runtime {
22+
/// <summary>
23+
/// PEP 525 async generator object — what an <c>async def</c> with <c>yield</c> returns.
24+
/// Façade over the DLR's <c>IAsyncEnumerable&lt;object&gt;</c> (produced by AsyncEnumerableExpression /
25+
/// AsyncHelpers.DriveAsyncEnumerable) plus the channels that let the consumer drive it:
26+
/// <list type="bullet">
27+
/// <item><c>_sendSlot</c> — the value delivered into the body at the next <c>yield</c> resume
28+
/// (<c>x = yield z</c>). Set by <see cref="asend"/> / <see cref="__anext__"/> before advancing.</item>
29+
/// <item><c>_throwSlot</c> — an exception rethrown at the next <c>yield</c> resume (athrow/aclose).</item>
30+
/// <item><c>_cts</c> — backs the underlying enumerator's cancellation.</item>
31+
/// </list>
32+
/// </summary>
33+
[PythonType("async_generator")]
34+
public sealed class PythonAsyncGenerator {
35+
private readonly IAsyncEnumerator<object?> _enumerator;
36+
private readonly StrongBox<object?> _sendSlot;
37+
private readonly StrongBox<Exception?> _throwSlot;
38+
private readonly CancellationTokenSource _cts;
39+
private readonly string _name;
40+
private bool _started;
41+
private bool _closed;
42+
43+
internal PythonAsyncGenerator(IAsyncEnumerable<object?> source,
44+
StrongBox<object?> sendSlot,
45+
StrongBox<Exception?> throwSlot,
46+
CancellationTokenSource cts,
47+
string name) {
48+
_enumerator = source.GetAsyncEnumerator(cts.Token);
49+
_sendSlot = sendSlot;
50+
_throwSlot = throwSlot;
51+
_cts = cts;
52+
_name = name;
53+
}
54+
55+
public PythonAsyncGenerator __aiter__() => this;
56+
57+
/// <summary>
58+
/// Advance the generator with no value sent in (equivalent to <c>asend(None)</c>). Returns an
59+
/// awaitable yielding the next produced value or raising StopAsyncIteration.
60+
/// </summary>
61+
public object __anext__() {
62+
_sendSlot.Value = null;
63+
return new AsyncEnumeratorAwaitable<object?>(_enumerator);
64+
}
65+
66+
/// <summary>
67+
/// Send a value into the generator; it becomes the value of the <c>yield</c> the body is suspended
68+
/// at (<c>x = yield z</c>). The first call after creation must send None (the body hasn't reached a
69+
/// yield yet), matching CPython.
70+
/// </summary>
71+
public object asend(object? value) {
72+
if (!_started) {
73+
_started = true;
74+
if (value is not null) {
75+
throw PythonOps.TypeError("can't send non-None value to a just-started async generator");
76+
}
77+
}
78+
_sendSlot.Value = value;
79+
return new AsyncEnumeratorAwaitable<object?>(_enumerator);
80+
}
81+
82+
[LightThrowing]
83+
public object athrow(object? type) => athrow(type, null, null);
84+
85+
[LightThrowing]
86+
public object athrow(object? type, object? value) => athrow(type, value, null);
87+
88+
/// <summary>
89+
/// Throw an exception into the generator at the suspended <c>yield</c>. Returns an awaitable; when
90+
/// awaited the exception is rethrown at the resume point — if the body catches it and yields again
91+
/// that value is produced, otherwise the exception propagates (or StopAsyncIteration if the body
92+
/// finishes).
93+
/// </summary>
94+
[LightThrowing]
95+
public object athrow(object? type, object? value, object? traceback) {
96+
// Validate shape (mirrors PythonGenerator/PythonCoroutine.throw).
97+
if (type is Exception || type is PythonExceptions.BaseException) {
98+
if (value is not null)
99+
return LightExceptions.Throw(PythonOps.TypeError("instance exception may not have a separate value"));
100+
} else if (type is PythonType pt && typeof(PythonExceptions.BaseException).IsAssignableFrom(pt.UnderlyingSystemType)) {
101+
// ok — class form
102+
} else {
103+
return LightExceptions.Throw(PythonOps.TypeError(
104+
"exceptions must be classes or instances deriving from BaseException, not {0}",
105+
PythonOps.GetPythonTypeName(type)));
106+
}
107+
108+
Exception ex = PythonOps.MakeExceptionForGenerator(DefaultContext.Default, type, value, traceback, cause: null);
109+
_started = true;
110+
_throwSlot.Value = ex;
111+
_sendSlot.Value = null;
112+
return new AsyncEnumeratorAwaitable<object?>(_enumerator);
113+
}
114+
115+
/// <summary>
116+
/// Close the generator: returns an awaitable that injects GeneratorExit at the suspended <c>yield</c>
117+
/// (running the body's try/finally), then disposes the underlying async iterator. Idempotent.
118+
/// </summary>
119+
public object aclose() {
120+
return AcloseAsync();
121+
}
122+
123+
private async Task<object?> AcloseAsync() {
124+
if (_closed) return null;
125+
_closed = true;
126+
if (_started) {
127+
_throwSlot.Value = new GeneratorExitException();
128+
_sendSlot.Value = null;
129+
try {
130+
bool produced = await _enumerator.MoveNextAsync();
131+
if (produced) {
132+
throw PythonOps.RuntimeError("async generator ignored GeneratorExit");
133+
}
134+
} catch (GeneratorExitException) {
135+
// expected — the body let GeneratorExit propagate
136+
} catch (StopIterationException) {
137+
// body finished — also fine
138+
}
139+
}
140+
await _enumerator.DisposeAsync();
141+
_cts.Cancel();
142+
return null;
143+
}
144+
145+
public string __name__ => _name;
146+
147+
public string __qualname__ => _name;
148+
149+
public bool ag_running => _started && !_closed;
150+
}
151+
}
152+
153+
#endif

0 commit comments

Comments
 (0)