Issue #28008: Implement PEP 530 -- asynchronous comprehensions.
diff --git a/Python/Python-ast.c b/Python/Python-ast.c
index 6ab57df..f10e315 100644
--- a/Python/Python-ast.c
+++ b/Python/Python-ast.c
@@ -435,10 +435,12 @@
static PyTypeObject *comprehension_type;
static PyObject* ast2obj_comprehension(void*);
_Py_IDENTIFIER(ifs);
+_Py_IDENTIFIER(is_async);
static char *comprehension_fields[]={
"target",
"iter",
"ifs",
+ "is_async",
};
static PyTypeObject *excepthandler_type;
static char *excepthandler_attributes[] = {
@@ -1148,7 +1150,7 @@
NotIn_singleton = PyType_GenericNew(NotIn_type, NULL, NULL);
if (!NotIn_singleton) return 0;
comprehension_type = make_type("comprehension", &AST_type,
- comprehension_fields, 3);
+ comprehension_fields, 4);
if (!comprehension_type) return 0;
if (!add_attributes(comprehension_type, NULL, 0)) return 0;
excepthandler_type = make_type("excepthandler", &AST_type, NULL, 0);
@@ -2445,7 +2447,8 @@
}
comprehension_ty
-comprehension(expr_ty target, expr_ty iter, asdl_seq * ifs, PyArena *arena)
+comprehension(expr_ty target, expr_ty iter, asdl_seq * ifs, int is_async,
+ PyArena *arena)
{
comprehension_ty p;
if (!target) {
@@ -2464,6 +2467,7 @@
p->target = target;
p->iter = iter;
p->ifs = ifs;
+ p->is_async = is_async;
return p;
}
@@ -3722,6 +3726,11 @@
if (_PyObject_SetAttrId(result, &PyId_ifs, value) == -1)
goto failed;
Py_DECREF(value);
+ value = ast2obj_int(o->is_async);
+ if (!value) goto failed;
+ if (_PyObject_SetAttrId(result, &PyId_is_async, value) == -1)
+ goto failed;
+ Py_DECREF(value);
return result;
failed:
Py_XDECREF(value);
@@ -7146,6 +7155,7 @@
expr_ty target;
expr_ty iter;
asdl_seq* ifs;
+ int is_async;
if (_PyObject_HasAttrId(obj, &PyId_target)) {
int res;
@@ -7193,7 +7203,18 @@
PyErr_SetString(PyExc_TypeError, "required field \"ifs\" missing from comprehension");
return 1;
}
- *out = comprehension(target, iter, ifs, arena);
+ if (_PyObject_HasAttrId(obj, &PyId_is_async)) {
+ int res;
+ tmp = _PyObject_GetAttrId(obj, &PyId_is_async);
+ if (tmp == NULL) goto failed;
+ res = obj2ast_int(tmp, &is_async, arena);
+ if (res != 0) goto failed;
+ Py_CLEAR(tmp);
+ } else {
+ PyErr_SetString(PyExc_TypeError, "required field \"is_async\" missing from comprehension");
+ return 1;
+ }
+ *out = comprehension(target, iter, ifs, is_async, arena);
return 0;
failed:
Py_XDECREF(tmp);
diff --git a/Python/ast.c b/Python/ast.c
index e89ec22..3719332 100644
--- a/Python/ast.c
+++ b/Python/ast.c
@@ -1747,14 +1747,21 @@
count_comp_fors(struct compiling *c, const node *n)
{
int n_fors = 0;
+ int is_async;
count_comp_for:
+ is_async = 0;
n_fors++;
REQ(n, comp_for);
- if (NCH(n) == 5)
- n = CHILD(n, 4);
- else
+ if (TYPE(CHILD(n, 0)) == ASYNC) {
+ is_async = 1;
+ }
+ if (NCH(n) == (5 + is_async)) {
+ n = CHILD(n, 4 + is_async);
+ }
+ else {
return n_fors;
+ }
count_comp_iter:
REQ(n, comp_iter);
n = CHILD(n, 0);
@@ -1817,14 +1824,19 @@
asdl_seq *t;
expr_ty expression, first;
node *for_ch;
+ int is_async = 0;
REQ(n, comp_for);
- for_ch = CHILD(n, 1);
+ if (TYPE(CHILD(n, 0)) == ASYNC) {
+ is_async = 1;
+ }
+
+ for_ch = CHILD(n, 1 + is_async);
t = ast_for_exprlist(c, for_ch, Store);
if (!t)
return NULL;
- expression = ast_for_expr(c, CHILD(n, 3));
+ expression = ast_for_expr(c, CHILD(n, 3 + is_async));
if (!expression)
return NULL;
@@ -1832,19 +1844,20 @@
(x for x, in ...) has 1 element in t, but still requires a Tuple. */
first = (expr_ty)asdl_seq_GET(t, 0);
if (NCH(for_ch) == 1)
- comp = comprehension(first, expression, NULL, c->c_arena);
+ comp = comprehension(first, expression, NULL,
+ is_async, c->c_arena);
else
- comp = comprehension(Tuple(t, Store, first->lineno, first->col_offset,
- c->c_arena),
- expression, NULL, c->c_arena);
+ comp = comprehension(Tuple(t, Store, first->lineno,
+ first->col_offset, c->c_arena),
+ expression, NULL, is_async, c->c_arena);
if (!comp)
return NULL;
- if (NCH(n) == 5) {
+ if (NCH(n) == (5 + is_async)) {
int j, n_ifs;
asdl_seq *ifs;
- n = CHILD(n, 4);
+ n = CHILD(n, 4 + is_async);
n_ifs = count_comp_ifs(c, n);
if (n_ifs == -1)
return NULL;
diff --git a/Python/compile.c b/Python/compile.c
index faae4f5..20fe4bc 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -202,6 +202,16 @@
static int compiler_try_except(struct compiler *, stmt_ty);
static int compiler_set_qualname(struct compiler *);
+static int compiler_sync_comprehension_generator(
+ struct compiler *c,
+ asdl_seq *generators, int gen_index,
+ expr_ty elt, expr_ty val, int type);
+
+static int compiler_async_comprehension_generator(
+ struct compiler *c,
+ asdl_seq *generators, int gen_index,
+ expr_ty elt, expr_ty val, int type);
+
static PyCodeObject *assemble(struct compiler *, int addNone);
static PyObject *__doc__;
@@ -2165,14 +2175,14 @@
static int
compiler_async_for(struct compiler *c, stmt_ty s)
{
- static PyObject *stopiter_error = NULL;
+ _Py_IDENTIFIER(StopAsyncIteration);
+
basicblock *try, *except, *end, *after_try, *try_cleanup,
*after_loop, *after_loop_else;
- if (stopiter_error == NULL) {
- stopiter_error = PyUnicode_InternFromString("StopAsyncIteration");
- if (stopiter_error == NULL)
- return 0;
+ PyObject *stop_aiter_error = _PyUnicode_FromId(&PyId_StopAsyncIteration);
+ if (stop_aiter_error == NULL) {
+ return 0;
}
try = compiler_new_block(c);
@@ -2214,7 +2224,7 @@
compiler_use_next_block(c, except);
ADDOP(c, DUP_TOP);
- ADDOP_O(c, LOAD_GLOBAL, stopiter_error, names);
+ ADDOP_O(c, LOAD_GLOBAL, stop_aiter_error, names);
ADDOP_I(c, COMPARE_OP, PyCmp_EXC_MATCH);
ADDOP_JABS(c, POP_JUMP_IF_FALSE, try_cleanup);
@@ -3627,11 +3637,28 @@
- iterate over the generator sequence instead of using recursion
*/
+
static int
compiler_comprehension_generator(struct compiler *c,
asdl_seq *generators, int gen_index,
expr_ty elt, expr_ty val, int type)
{
+ comprehension_ty gen;
+ gen = (comprehension_ty)asdl_seq_GET(generators, gen_index);
+ if (gen->is_async) {
+ return compiler_async_comprehension_generator(
+ c, generators, gen_index, elt, val, type);
+ } else {
+ return compiler_sync_comprehension_generator(
+ c, generators, gen_index, elt, val, type);
+ }
+}
+
+static int
+compiler_sync_comprehension_generator(struct compiler *c,
+ asdl_seq *generators, int gen_index,
+ expr_ty elt, expr_ty val, int type)
+{
/* generate code for the iterator, then each of the ifs,
and then write to the element */
@@ -3718,20 +3745,167 @@
}
static int
+compiler_async_comprehension_generator(struct compiler *c,
+ asdl_seq *generators, int gen_index,
+ expr_ty elt, expr_ty val, int type)
+{
+ _Py_IDENTIFIER(StopAsyncIteration);
+
+ comprehension_ty gen;
+ basicblock *anchor, *skip, *if_cleanup, *try,
+ *after_try, *except, *try_cleanup;
+ Py_ssize_t i, n;
+
+ PyObject *stop_aiter_error = _PyUnicode_FromId(&PyId_StopAsyncIteration);
+ if (stop_aiter_error == NULL) {
+ return 0;
+ }
+
+ try = compiler_new_block(c);
+ after_try = compiler_new_block(c);
+ try_cleanup = compiler_new_block(c);
+ except = compiler_new_block(c);
+ skip = compiler_new_block(c);
+ if_cleanup = compiler_new_block(c);
+ anchor = compiler_new_block(c);
+
+ if (skip == NULL || if_cleanup == NULL || anchor == NULL ||
+ try == NULL || after_try == NULL ||
+ except == NULL || after_try == NULL) {
+ return 0;
+ }
+
+ gen = (comprehension_ty)asdl_seq_GET(generators, gen_index);
+
+ if (gen_index == 0) {
+ /* Receive outermost iter as an implicit argument */
+ c->u->u_argcount = 1;
+ ADDOP_I(c, LOAD_FAST, 0);
+ }
+ else {
+ /* Sub-iter - calculate on the fly */
+ VISIT(c, expr, gen->iter);
+ ADDOP(c, GET_AITER);
+ ADDOP_O(c, LOAD_CONST, Py_None, consts);
+ ADDOP(c, YIELD_FROM);
+ }
+
+ compiler_use_next_block(c, try);
+
+
+ ADDOP_JREL(c, SETUP_EXCEPT, except);
+ if (!compiler_push_fblock(c, EXCEPT, try))
+ return 0;
+
+ ADDOP(c, GET_ANEXT);
+ ADDOP_O(c, LOAD_CONST, Py_None, consts);
+ ADDOP(c, YIELD_FROM);
+ VISIT(c, expr, gen->target);
+ ADDOP(c, POP_BLOCK);
+ compiler_pop_fblock(c, EXCEPT, try);
+ ADDOP_JREL(c, JUMP_FORWARD, after_try);
+
+
+ compiler_use_next_block(c, except);
+ ADDOP(c, DUP_TOP);
+ ADDOP_O(c, LOAD_GLOBAL, stop_aiter_error, names);
+ ADDOP_I(c, COMPARE_OP, PyCmp_EXC_MATCH);
+ ADDOP_JABS(c, POP_JUMP_IF_FALSE, try_cleanup);
+
+ ADDOP(c, POP_TOP);
+ ADDOP(c, POP_TOP);
+ ADDOP(c, POP_TOP);
+ ADDOP(c, POP_EXCEPT); /* for SETUP_EXCEPT */
+ ADDOP_JABS(c, JUMP_ABSOLUTE, anchor);
+
+
+ compiler_use_next_block(c, try_cleanup);
+ ADDOP(c, END_FINALLY);
+
+ compiler_use_next_block(c, after_try);
+
+ n = asdl_seq_LEN(gen->ifs);
+ for (i = 0; i < n; i++) {
+ expr_ty e = (expr_ty)asdl_seq_GET(gen->ifs, i);
+ VISIT(c, expr, e);
+ ADDOP_JABS(c, POP_JUMP_IF_FALSE, if_cleanup);
+ NEXT_BLOCK(c);
+ }
+
+ if (++gen_index < asdl_seq_LEN(generators))
+ if (!compiler_comprehension_generator(c,
+ generators, gen_index,
+ elt, val, type))
+ return 0;
+
+ /* only append after the last for generator */
+ if (gen_index >= asdl_seq_LEN(generators)) {
+ /* comprehension specific code */
+ switch (type) {
+ case COMP_GENEXP:
+ VISIT(c, expr, elt);
+ ADDOP(c, YIELD_VALUE);
+ ADDOP(c, POP_TOP);
+ break;
+ case COMP_LISTCOMP:
+ VISIT(c, expr, elt);
+ ADDOP_I(c, LIST_APPEND, gen_index + 1);
+ break;
+ case COMP_SETCOMP:
+ VISIT(c, expr, elt);
+ ADDOP_I(c, SET_ADD, gen_index + 1);
+ break;
+ case COMP_DICTCOMP:
+ /* With 'd[k] = v', v is evaluated before k, so we do
+ the same. */
+ VISIT(c, expr, val);
+ VISIT(c, expr, elt);
+ ADDOP_I(c, MAP_ADD, gen_index + 1);
+ break;
+ default:
+ return 0;
+ }
+
+ compiler_use_next_block(c, skip);
+ }
+ compiler_use_next_block(c, if_cleanup);
+ ADDOP_JABS(c, JUMP_ABSOLUTE, try);
+ compiler_use_next_block(c, anchor);
+ ADDOP(c, POP_TOP);
+
+ return 1;
+}
+
+static int
compiler_comprehension(struct compiler *c, expr_ty e, int type,
identifier name, asdl_seq *generators, expr_ty elt,
expr_ty val)
{
PyCodeObject *co = NULL;
- expr_ty outermost_iter;
+ comprehension_ty outermost;
PyObject *qualname = NULL;
+ int is_async_function = c->u->u_ste->ste_coroutine;
+ int is_async_generator = 0;
- outermost_iter = ((comprehension_ty)
- asdl_seq_GET(generators, 0))->iter;
+ outermost = (comprehension_ty) asdl_seq_GET(generators, 0);
if (!compiler_enter_scope(c, name, COMPILER_SCOPE_COMPREHENSION,
(void *)e, e->lineno))
+ {
goto error;
+ }
+
+ is_async_generator = c->u->u_ste->ste_coroutine;
+
+ if (is_async_generator && !is_async_function) {
+ if (e->lineno > c->u->u_lineno) {
+ c->u->u_lineno = e->lineno;
+ c->u->u_lineno_set = 0;
+ }
+ compiler_error(c, "asynchronous comprehension outside of "
+ "an asynchronous function");
+ goto error_in_scope;
+ }
if (type != COMP_GENEXP) {
int op;
@@ -3774,9 +3948,24 @@
Py_DECREF(qualname);
Py_DECREF(co);
- VISIT(c, expr, outermost_iter);
- ADDOP(c, GET_ITER);
+ VISIT(c, expr, outermost->iter);
+
+ if (outermost->is_async) {
+ ADDOP(c, GET_AITER);
+ ADDOP_O(c, LOAD_CONST, Py_None, consts);
+ ADDOP(c, YIELD_FROM);
+ } else {
+ ADDOP(c, GET_ITER);
+ }
+
ADDOP_I(c, CALL_FUNCTION, 1);
+
+ if (is_async_generator && type != COMP_GENEXP) {
+ ADDOP(c, GET_AWAITABLE);
+ ADDOP_O(c, LOAD_CONST, Py_None, consts);
+ ADDOP(c, YIELD_FROM);
+ }
+
return 1;
error_in_scope:
compiler_exit_scope(c);
@@ -4140,11 +4329,8 @@
if (c->u->u_ste->ste_type != FunctionBlock)
return compiler_error(c, "'await' outside function");
- if (c->u->u_scope_type == COMPILER_SCOPE_COMPREHENSION)
- return compiler_error(
- c, "'await' expressions in comprehensions are not supported");
-
- if (c->u->u_scope_type != COMPILER_SCOPE_ASYNC_FUNCTION)
+ if (c->u->u_scope_type != COMPILER_SCOPE_ASYNC_FUNCTION &&
+ c->u->u_scope_type != COMPILER_SCOPE_COMPREHENSION)
return compiler_error(c, "'await' outside async function");
VISIT(c, expr, e->v.Await.value);
diff --git a/Python/graminit.c b/Python/graminit.c
index c11b831..f2584e0 100644
--- a/Python/graminit.c
+++ b/Python/graminit.c
@@ -1812,32 +1812,37 @@
{2, arcs_80_0},
{1, arcs_80_1},
};
-static arc arcs_81_0[1] = {
- {101, 1},
+static arc arcs_81_0[2] = {
+ {21, 1},
+ {101, 2},
};
static arc arcs_81_1[1] = {
- {66, 2},
+ {101, 2},
};
static arc arcs_81_2[1] = {
- {102, 3},
+ {66, 3},
};
static arc arcs_81_3[1] = {
- {112, 4},
+ {102, 4},
};
-static arc arcs_81_4[2] = {
- {171, 5},
- {0, 4},
+static arc arcs_81_4[1] = {
+ {112, 5},
};
-static arc arcs_81_5[1] = {
+static arc arcs_81_5[2] = {
+ {171, 6},
{0, 5},
};
-static state states_81[6] = {
- {1, arcs_81_0},
+static arc arcs_81_6[1] = {
+ {0, 6},
+};
+static state states_81[7] = {
+ {2, arcs_81_0},
{1, arcs_81_1},
{1, arcs_81_2},
{1, arcs_81_3},
- {2, arcs_81_4},
- {1, arcs_81_5},
+ {1, arcs_81_4},
+ {2, arcs_81_5},
+ {1, arcs_81_6},
};
static arc arcs_82_0[1] = {
{97, 1},
@@ -2060,9 +2065,9 @@
{335, "argument", 0, 4, states_79,
"\000\040\200\000\006\000\000\000\000\000\010\000\000\000\020\002\000\300\220\050\037\000"},
{336, "comp_iter", 0, 2, states_80,
- "\000\000\000\000\000\000\000\000\000\000\000\000\042\000\000\000\000\000\000\000\000\000"},
- {337, "comp_for", 0, 6, states_81,
- "\000\000\000\000\000\000\000\000\000\000\000\000\040\000\000\000\000\000\000\000\000\000"},
+ "\000\000\040\000\000\000\000\000\000\000\000\000\042\000\000\000\000\000\000\000\000\000"},
+ {337, "comp_for", 0, 7, states_81,
+ "\000\000\040\000\000\000\000\000\000\000\000\000\040\000\000\000\000\000\000\000\000\000"},
{338, "comp_if", 0, 4, states_82,
"\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000"},
{339, "encoding_decl", 0, 2, states_83,
diff --git a/Python/symtable.c b/Python/symtable.c
index 9325dc1..f762904 100644
--- a/Python/symtable.c
+++ b/Python/symtable.c
@@ -1682,6 +1682,9 @@
VISIT(st, expr, lc->target);
VISIT(st, expr, lc->iter);
VISIT_SEQ(st, expr, lc->ifs);
+ if (lc->is_async) {
+ st->st_cur->ste_coroutine = 1;
+ }
return 1;
}
@@ -1734,6 +1737,9 @@
return 0;
}
st->st_cur->ste_generator = is_generator;
+ if (outermost->is_async) {
+ st->st_cur->ste_coroutine = 1;
+ }
/* Outermost iter is received as an argument */
if (!symtable_implicit_arg(st, 0)) {
symtable_exit_block(st, (void *)e);