reflect with statements with multiple items in the AST (closes #12106)

7 files changed, 182 insertions, 85 deletions

diff --git a/Include/Python-ast.h b/Include/Python-ast.h
index 0ad788b..74c5264 100644
--- a/Include/Python-ast.h
+++ b/Include/Python-ast.h
@@ -36,6 +36,8 @@ typedef struct _keyword *keyword_ty;
 
 typedef struct _alias *alias_ty;
 
+typedef struct _withitem *withitem_ty;
+
 
 enum _mod_kind {Module_kind=1, Interactive_kind=2, Expression_kind=3,
                  Suite_kind=4};
@@ -128,8 +130,7 @@ struct _stmt {
                 } If;
                 
                 struct {
-                        expr_ty context_expr;
-                        expr_ty optional_vars;
+                        asdl_seq *items;
                         asdl_seq *body;
                 } With;
                 
@@ -383,6 +384,11 @@ struct _alias {
         identifier asname;
 };
 
+struct _withitem {
+        expr_ty context_expr;
+        expr_ty optional_vars;
+};
+
 
 #define Module(a0, a1) _Py_Module(a0, a1)
 mod_ty _Py_Module(asdl_seq * body, PyArena *arena);
@@ -421,9 +427,9 @@ stmt_ty _Py_While(expr_ty test, asdl_seq * body, asdl_seq * orelse, int lineno,
 #define If(a0, a1, a2, a3, a4, a5) _Py_If(a0, a1, a2, a3, a4, a5)
 stmt_ty _Py_If(expr_ty test, asdl_seq * body, asdl_seq * orelse, int lineno,
                int col_offset, PyArena *arena);
-#define With(a0, a1, a2, a3, a4, a5) _Py_With(a0, a1, a2, a3, a4, a5)
-stmt_ty _Py_With(expr_ty context_expr, expr_ty optional_vars, asdl_seq * body,
-                 int lineno, int col_offset, PyArena *arena);
+#define With(a0, a1, a2, a3, a4) _Py_With(a0, a1, a2, a3, a4)
+stmt_ty _Py_With(asdl_seq * items, asdl_seq * body, int lineno, int col_offset,
+                 PyArena *arena);
 #define Raise(a0, a1, a2, a3, a4) _Py_Raise(a0, a1, a2, a3, a4)
 stmt_ty _Py_Raise(expr_ty exc, expr_ty cause, int lineno, int col_offset,
                   PyArena *arena);
@@ -547,6 +553,9 @@ arg_ty _Py_arg(identifier arg, expr_ty annotation, PyArena *arena);
 keyword_ty _Py_keyword(identifier arg, expr_ty value, PyArena *arena);
 #define alias(a0, a1, a2) _Py_alias(a0, a1, a2)
 alias_ty _Py_alias(identifier name, identifier asname, PyArena *arena);
+#define withitem(a0, a1, a2) _Py_withitem(a0, a1, a2)
+withitem_ty _Py_withitem(expr_ty context_expr, expr_ty optional_vars, PyArena
+                         *arena);
 
 PyObject* PyAST_mod2obj(mod_ty t);
 mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena, int mode);
diff --git a/Misc/NEWS b/Misc/NEWS
index 2a916fd..9c00590 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -10,6 +10,9 @@ What's New in Python 3.3 Alpha 1?
 Core and Builtins
 -----------------
 
+- Issue #12106: The use of the multiple-with shorthand syntax is now reflected
+  in the AST.
+
 - Issue #12190: Try to use the same filename object when compiling unmarshalling
   a code objects in the same file.
 
diff --git a/Parser/Python.asdl b/Parser/Python.asdl
index 8e2e1ac..de48643 100644
--- a/Parser/Python.asdl
+++ b/Parser/Python.asdl
@@ -28,7 +28,7 @@ module Python
 	      | For(expr target, expr iter, stmt* body, stmt* orelse)
 	      | While(expr test, stmt* body, stmt* orelse)
 	      | If(expr test, stmt* body, stmt* orelse)
-	      | With(expr context_expr, expr? optional_vars, stmt* body)
+	      | With(withitem* items, stmt* body)
 
 	      | Raise(expr? exc, expr? cause)
 	      | TryExcept(stmt* body, excepthandler* handlers, stmt* orelse)
@@ -115,5 +115,7 @@ module Python
 
         -- import name with optional 'as' alias.
         alias = (identifier name, identifier? asname)
+
+        withitem = (expr context_expr, expr? optional_vars)
 }
 
diff --git a/Python/Python-ast.c b/Python/Python-ast.c
index 6b1ea3c..2364db3 100644
--- a/Python/Python-ast.c
+++ b/Python/Python-ast.c
@@ -95,8 +95,7 @@ static char *If_fields[]={
 };
 static PyTypeObject *With_type;
 static char *With_fields[]={
-        "context_expr",
-        "optional_vars",
+        "items",
         "body",
 };
 static PyTypeObject *Raise_type;
@@ -392,6 +391,12 @@ static char *alias_fields[]={
         "name",
         "asname",
 };
+static PyTypeObject *withitem_type;
+static PyObject* ast2obj_withitem(void*);
+static char *withitem_fields[]={
+        "context_expr",
+        "optional_vars",
+};
 
 
 static int
@@ -680,7 +685,7 @@ static int init_types(void)
         if (!While_type) return 0;
         If_type = make_type("If", stmt_type, If_fields, 3);
         if (!If_type) return 0;
-        With_type = make_type("With", stmt_type, With_fields, 3);
+        With_type = make_type("With", stmt_type, With_fields, 2);
         if (!With_type) return 0;
         Raise_type = make_type("Raise", stmt_type, Raise_fields, 2);
         if (!Raise_type) return 0;
@@ -938,6 +943,8 @@ static int init_types(void)
         if (!keyword_type) return 0;
         alias_type = make_type("alias", &AST_type, alias_fields, 2);
         if (!alias_type) return 0;
+        withitem_type = make_type("withitem", &AST_type, withitem_fields, 2);
+        if (!withitem_type) return 0;
         initialized = 1;
         return 1;
 }
@@ -960,6 +967,7 @@ static int obj2ast_arguments(PyObject* obj, arguments_ty* out, PyArena* arena);
 static int obj2ast_arg(PyObject* obj, arg_ty* out, PyArena* arena);
 static int obj2ast_keyword(PyObject* obj, keyword_ty* out, PyArena* arena);
 static int obj2ast_alias(PyObject* obj, alias_ty* out, PyArena* arena);
+static int obj2ast_withitem(PyObject* obj, withitem_ty* out, PyArena* arena);
 
 mod_ty
 Module(asdl_seq * body, PyArena *arena)
@@ -1225,21 +1233,15 @@ If(expr_ty test, asdl_seq * body, asdl_seq * orelse, int lineno, int
 }
 
 stmt_ty
-With(expr_ty context_expr, expr_ty optional_vars, asdl_seq * body, int lineno,
-     int col_offset, PyArena *arena)
+With(asdl_seq * items, asdl_seq * body, int lineno, int col_offset, PyArena
+     *arena)
 {
         stmt_ty p;
-        if (!context_expr) {
-                PyErr_SetString(PyExc_ValueError,
-                                "field context_expr is required for With");
-                return NULL;
-        }
         p = (stmt_ty)PyArena_Malloc(arena, sizeof(*p));
         if (!p)
                 return NULL;
         p->kind = With_kind;
-        p->v.With.context_expr = context_expr;
-        p->v.With.optional_vars = optional_vars;
+        p->v.With.items = items;
         p->v.With.body = body;
         p->lineno = lineno;
         p->col_offset = col_offset;
@@ -2135,6 +2137,23 @@ alias(identifier name, identifier asname, PyArena *arena)
         return p;
 }
 
+withitem_ty
+withitem(expr_ty context_expr, expr_ty optional_vars, PyArena *arena)
+{
+        withitem_ty p;
+        if (!context_expr) {
+                PyErr_SetString(PyExc_ValueError,
+                                "field context_expr is required for withitem");
+                return NULL;
+        }
+        p = (withitem_ty)PyArena_Malloc(arena, sizeof(*p));
+        if (!p)
+                return NULL;
+        p->context_expr = context_expr;
+        p->optional_vars = optional_vars;
+        return p;
+}
+
 
 PyObject*
 ast2obj_mod(void* _o)
@@ -2390,15 +2409,9 @@ ast2obj_stmt(void* _o)
         case With_kind:
                 result = PyType_GenericNew(With_type, NULL, NULL);
                 if (!result) goto failed;
-                value = ast2obj_expr(o->v.With.context_expr);
+                value = ast2obj_list(o->v.With.items, ast2obj_withitem);
                 if (!value) goto failed;
-                if (PyObject_SetAttrString(result, "context_expr", value) == -1)
-                        goto failed;
-                Py_DECREF(value);
-                value = ast2obj_expr(o->v.With.optional_vars);
-                if (!value) goto failed;
-                if (PyObject_SetAttrString(result, "optional_vars", value) ==
-                    -1)
+                if (PyObject_SetAttrString(result, "items", value) == -1)
                         goto failed;
                 Py_DECREF(value);
                 value = ast2obj_list(o->v.With.body, ast2obj_stmt);
@@ -3370,6 +3383,35 @@ failed:
         return NULL;
 }
 
+PyObject*
+ast2obj_withitem(void* _o)
+{
+        withitem_ty o = (withitem_ty)_o;
+        PyObject *result = NULL, *value = NULL;
+        if (!o) {
+                Py_INCREF(Py_None);
+                return Py_None;
+        }
+
+        result = PyType_GenericNew(withitem_type, NULL, NULL);
+        if (!result) return NULL;
+        value = ast2obj_expr(o->context_expr);
+        if (!value) goto failed;
+        if (PyObject_SetAttrString(result, "context_expr", value) == -1)
+                goto failed;
+        Py_DECREF(value);
+        value = ast2obj_expr(o->optional_vars);
+        if (!value) goto failed;
+        if (PyObject_SetAttrString(result, "optional_vars", value) == -1)
+                goto failed;
+        Py_DECREF(value);
+        return result;
+failed:
+        Py_XDECREF(value);
+        Py_XDECREF(result);
+        return NULL;
+}
+
 
 int
 obj2ast_mod(PyObject* obj, mod_ty* out, PyArena* arena)
@@ -4210,33 +4252,34 @@ obj2ast_stmt(PyObject* obj, stmt_ty* out, PyArena* arena)
                 return 1;
         }
         if (isinstance) {
-                expr_ty context_expr;
-                expr_ty optional_vars;
+                asdl_seq* items;
                 asdl_seq* body;
 
-                if (PyObject_HasAttrString(obj, "context_expr")) {
+                if (PyObject_HasAttrString(obj, "items")) {
                         int res;
-                        tmp = PyObject_GetAttrString(obj, "context_expr");
+                        Py_ssize_t len;
+                        Py_ssize_t i;
+                        tmp = PyObject_GetAttrString(obj, "items");
                         if (tmp == NULL) goto failed;
-                        res = obj2ast_expr(tmp, &context_expr, arena);
-                        if (res != 0) goto failed;
+                        if (!PyList_Check(tmp)) {
+                                PyErr_Format(PyExc_TypeError, "With field \"items\" must be a list, not a %.200s", tmp->ob_type->tp_name);
+                                goto failed;
+                        }
+                        len = PyList_GET_SIZE(tmp);
+                        items = asdl_seq_new(len, arena);
+                        if (items == NULL) goto failed;
+                        for (i = 0; i < len; i++) {
+                                withitem_ty value;
+                                res = obj2ast_withitem(PyList_GET_ITEM(tmp, i), &value, arena);
+                                if (res != 0) goto failed;
+                                asdl_seq_SET(items, i, value);
+                        }
                         Py_XDECREF(tmp);
                         tmp = NULL;
                 } else {
-                        PyErr_SetString(PyExc_TypeError, "required field \"context_expr\" missing from With");
+                        PyErr_SetString(PyExc_TypeError, "required field \"items\" missing from With");
                         return 1;
                 }
-                if (PyObject_HasAttrString(obj, "optional_vars")) {
-                        int res;
-                        tmp = PyObject_GetAttrString(obj, "optional_vars");
-                        if (tmp == NULL) goto failed;
-                        res = obj2ast_expr(tmp, &optional_vars, arena);
-                        if (res != 0) goto failed;
-                        Py_XDECREF(tmp);
-                        tmp = NULL;
-                } else {
-                        optional_vars = NULL;
-                }
                 if (PyObject_HasAttrString(obj, "body")) {
                         int res;
                         Py_ssize_t len;
@@ -4262,8 +4305,7 @@ obj2ast_stmt(PyObject* obj, stmt_ty* out, PyArena* arena)
                         PyErr_SetString(PyExc_TypeError, "required field \"body\" missing from With");
                         return 1;
                 }
-                *out = With(context_expr, optional_vars, body, lineno,
-                            col_offset, arena);
+                *out = With(items, body, lineno, col_offset, arena);
                 if (*out == NULL) goto failed;
                 return 0;
         }
@@ -6723,6 +6765,43 @@ failed:
         return 1;
 }
 
+int
+obj2ast_withitem(PyObject* obj, withitem_ty* out, PyArena* arena)
+{
+        PyObject* tmp = NULL;
+        expr_ty context_expr;
+        expr_ty optional_vars;
+
+        if (PyObject_HasAttrString(obj, "context_expr")) {
+                int res;
+                tmp = PyObject_GetAttrString(obj, "context_expr");
+                if (tmp == NULL) goto failed;
+                res = obj2ast_expr(tmp, &context_expr, arena);
+                if (res != 0) goto failed;
+                Py_XDECREF(tmp);
+                tmp = NULL;
+        } else {
+                PyErr_SetString(PyExc_TypeError, "required field \"context_expr\" missing from withitem");
+                return 1;
+        }
+        if (PyObject_HasAttrString(obj, "optional_vars")) {
+                int res;
+                tmp = PyObject_GetAttrString(obj, "optional_vars");
+                if (tmp == NULL) goto failed;
+                res = obj2ast_expr(tmp, &optional_vars, arena);
+                if (res != 0) goto failed;
+                Py_XDECREF(tmp);
+                tmp = NULL;
+        } else {
+                optional_vars = NULL;
+        }
+        *out = withitem(context_expr, optional_vars, arena);
+        return 0;
+failed:
+        Py_XDECREF(tmp);
+        return 1;
+}
+
 
 static struct PyModuleDef _astmodule = {
   PyModuleDef_HEAD_INIT, "_ast"
@@ -6940,6 +7019,8 @@ PyInit__ast(void)
             return NULL;
         if (PyDict_SetItemString(d, "alias", (PyObject*)alias_type) < 0) return
             NULL;
+        if (PyDict_SetItemString(d, "withitem", (PyObject*)withitem_type) < 0)
+            return NULL;
         return m;
 }
 
diff --git a/Python/ast.c b/Python/ast.c
index 5b12da8..882452b 100644
--- a/Python/ast.c
+++ b/Python/ast.c
@@ -2967,8 +2967,8 @@ ast_for_try_stmt(struct compiling *c, const node *n)
 }
 
 /* with_item: test ['as' expr] */
-static stmt_ty
-ast_for_with_item(struct compiling *c, const node *n, asdl_seq *content)
+static withitem_ty
+ast_for_with_item(struct compiling *c, const node *n)
 {
     expr_ty context_expr, optional_vars = NULL;
 
@@ -2987,43 +2987,32 @@ ast_for_with_item(struct compiling *c, const node *n, asdl_seq *content)
         }
     }
 
-    return With(context_expr, optional_vars, content, LINENO(n),
-                n->n_col_offset, c->c_arena);
+    return withitem(context_expr, optional_vars, c->c_arena);
 }
 
 /* with_stmt: 'with' with_item (',' with_item)* ':' suite */
 static stmt_ty
 ast_for_with_stmt(struct compiling *c, const node *n)
 {
-    int i;
-    stmt_ty ret;
-    asdl_seq *inner;
+    int i, n_items;
+    asdl_seq *items, *body;
 
     REQ(n, with_stmt);
 
-    /* process the with items inside-out */
-    i = NCH(n) - 1;
-    /* the suite of the innermost with item is the suite of the with stmt */
-    inner = ast_for_suite(c, CHILD(n, i));
-    if (!inner)
-        return NULL;
-
-    for (;;) {
-        i -= 2;
-        ret = ast_for_with_item(c, CHILD(n, i), inner);
-        if (!ret)
+    n_items = (NCH(n) - 2) / 2;
+    items = asdl_seq_new(n_items, c->c_arena);
+    for (i = 1; i < NCH(n) - 2; i += 2) {
+        withitem_ty item = ast_for_with_item(c, CHILD(n, i));
+        if (!item)
             return NULL;
-        /* was this the last item? */
-        if (i == 1)
-            break;
-        /* if not, wrap the result so far in a new sequence */
-        inner = asdl_seq_new(1, c->c_arena);
-        if (!inner)
-            return NULL;
-        asdl_seq_SET(inner, 0, ret);
+        asdl_seq_SET(items, (i - 1) / 2, item);
     }
 
-    return ret;
+    body = ast_for_suite(c, CHILD(n, NCH(n) - 1));
+    if (!body)
+        return NULL;
+
+    return With(items, body, LINENO(n), n->n_col_offset, c->c_arena);
 }
 
 static stmt_ty
diff --git a/Python/compile.c b/Python/compile.c
index 34e7603..09983d2 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -179,7 +179,7 @@ static int compiler_in_loop(struct compiler *);
 static int inplace_binop(struct compiler *, operator_ty);
 static int expr_constant(struct compiler *, expr_ty);
 
-static int compiler_with(struct compiler *, stmt_ty);
+static int compiler_with(struct compiler *, stmt_ty, int);
 static int compiler_call_helper(struct compiler *c, int n,
                                 asdl_seq *args,
                                 asdl_seq *keywords,
@@ -2341,7 +2341,7 @@ compiler_visit_stmt(struct compiler *c, stmt_ty s)
     case Continue_kind:
         return compiler_continue(c);
     case With_kind:
-        return compiler_with(c, s);
+        return compiler_with(c, s, 0);
     }
     return 1;
 }
@@ -3068,9 +3068,10 @@ expr_constant(struct compiler *c, expr_ty e)
        exit(*exc)
  */
 static int
-compiler_with(struct compiler *c, stmt_ty s)
+compiler_with(struct compiler *c, stmt_ty s, int pos)
 {
     basicblock *block, *finally;
+    withitem_ty item = asdl_seq_GET(s->v.With.items, pos);
 
     assert(s->kind == With_kind);
 
@@ -3080,7 +3081,7 @@ compiler_with(struct compiler *c, stmt_ty s)
         return 0;
 
     /* Evaluate EXPR */
-    VISIT(c, expr, s->v.With.context_expr);
+    VISIT(c, expr, item->context_expr);
     ADDOP_JREL(c, SETUP_WITH, finally);
 
     /* SETUP_WITH pushes a finally block. */
@@ -3089,16 +3090,20 @@ compiler_with(struct compiler *c, stmt_ty s)
         return 0;
     }
 
-    if (s->v.With.optional_vars) {
-        VISIT(c, expr, s->v.With.optional_vars);
+    if (item->optional_vars) {
+        VISIT(c, expr, item->optional_vars);
     }
     else {
     /* Discard result from context.__enter__() */
         ADDOP(c, POP_TOP);
     }
 
-    /* BLOCK code */
-    VISIT_SEQ(c, stmt, s->v.With.body);
+    pos++;
+    if (pos == asdl_seq_LEN(s->v.With.items))
+        /* BLOCK code */
+        VISIT_SEQ(c, stmt, s->v.With.body)
+    else if (!compiler_with(c, s, pos))
+            return 0;
 
     /* End of try block; start the finally block */
     ADDOP(c, POP_BLOCK);
diff --git a/Python/symtable.c b/Python/symtable.c
index 8040665..d276254 100644
--- a/Python/symtable.c
+++ b/Python/symtable.c
@@ -185,6 +185,7 @@ static int symtable_visit_params(struct symtable *st, asdl_seq *args);
 static int symtable_visit_argannotations(struct symtable *st, asdl_seq *args);
 static int symtable_implicit_arg(struct symtable *st, int pos);
 static int symtable_visit_annotations(struct symtable *st, stmt_ty s);
+static int symtable_visit_withitem(struct symtable *st, withitem_ty item);
 
 
 static identifier top = NULL, lambda = NULL, genexpr = NULL,
@@ -1305,10 +1306,7 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
         /* nothing to do here */
         break;
     case With_kind:
-        VISIT(st, expr, s->v.With.context_expr);
-        if (s->v.With.optional_vars) {
-            VISIT(st, expr, s->v.With.optional_vars);
-        }
+        VISIT_SEQ(st, withitem, s->v.With.items);
         VISIT_SEQ(st, stmt, s->v.With.body);
         break;
     }
@@ -1540,6 +1538,16 @@ symtable_visit_excepthandler(struct symtable *st, excepthandler_ty eh)
     return 1;
 }
 
+static int
+symtable_visit_withitem(struct symtable *st, withitem_ty item)
+{
+    VISIT(st, expr, item->context_expr);
+    if (item->optional_vars) {
+        VISIT(st, expr, item->optional_vars);
+    }
+    return 1;
+}
+
 
 static int
 symtable_visit_alias(struct symtable *st, alias_ty a)