1 files changed, 83 insertions, 9 deletions

diff --git a/Python/symtable.c b/Python/symtable.c
index 677b604343..879e19ab79 100644
--- a/Python/symtable.c
+++ b/Python/symtable.c
@@ -31,6 +31,9 @@
 
 #define IMPORT_STAR_WARNING "import * only allowed at module level"
 
+#define NAMED_EXPR_COMP_IN_CLASS \
+"named expression within a comprehension cannot be used in a class body"
+
 static PySTEntryObject *
 ste_new(struct symtable *st, identifier name, _Py_block_ty block,
         void *key, int lineno, int col_offset)
@@ -75,6 +78,7 @@ ste_new(struct symtable *st, identifier name, _Py_block_ty block,
     ste->ste_child_free = 0;
     ste->ste_generator = 0;
     ste->ste_coroutine = 0;
+    ste->ste_comprehension = 0;
     ste->ste_returns_value = 0;
     ste->ste_needs_class_closure = 0;
 
@@ -972,7 +976,7 @@ symtable_lookup(struct symtable *st, PyObject *name)
 }
 
 static int
-symtable_add_def(struct symtable *st, PyObject *name, int flag)
+symtable_add_def_helper(struct symtable *st, PyObject *name, int flag, struct _symtable_entry *ste)
 {
     PyObject *o;
     PyObject *dict;
@@ -982,15 +986,15 @@ symtable_add_def(struct symtable *st, PyObject *name, int flag)
 
     if (!mangled)
         return 0;
-    dict = st->st_cur->ste_symbols;
+    dict = ste->ste_symbols;
     if ((o = PyDict_GetItem(dict, mangled))) {
         val = PyLong_AS_LONG(o);
         if ((flag & DEF_PARAM) && (val & DEF_PARAM)) {
             /* Is it better to use 'mangled' or 'name' here? */
             PyErr_Format(PyExc_SyntaxError, DUPLICATE_ARGUMENT, name);
             PyErr_SyntaxLocationObject(st->st_filename,
-                                       st->st_cur->ste_lineno,
-                                       st->st_cur->ste_col_offset + 1);
+                                       ste->ste_lineno,
+                                       ste->ste_col_offset + 1);
             goto error;
         }
         val |= flag;
@@ -1006,7 +1010,7 @@ symtable_add_def(struct symtable *st, PyObject *name, int flag)
     Py_DECREF(o);
 
     if (flag & DEF_PARAM) {
-        if (PyList_Append(st->st_cur->ste_varnames, mangled) < 0)
+        if (PyList_Append(ste->ste_varnames, mangled) < 0)
             goto error;
     } else      if (flag & DEF_GLOBAL) {
         /* XXX need to update DEF_GLOBAL for other flags too;
@@ -1032,6 +1036,11 @@ error:
     return 0;
 }
 
+static int
+symtable_add_def(struct symtable *st, PyObject *name, int flag) {
+    return symtable_add_def_helper(st, name, flag, st->st_cur);
+}
+
 /* VISIT, VISIT_SEQ and VIST_SEQ_TAIL take an ASDL type as their second argument.
    They use the ASDL name to synthesize the name of the C type and the visit
    function.
@@ -1082,7 +1091,7 @@ error:
 }
 
 static int
-symtable_record_directive(struct symtable *st, identifier name, stmt_ty s)
+symtable_record_directive(struct symtable *st, identifier name, int lineno, int col_offset)
 {
     PyObject *data, *mangled;
     int res;
@@ -1094,7 +1103,7 @@ symtable_record_directive(struct symtable *st, identifier name, stmt_ty s)
     mangled = _Py_Mangle(st->st_private, name);
     if (!mangled)
         return 0;
-    data = Py_BuildValue("(Nii)", mangled, s->lineno, s->col_offset);
+    data = Py_BuildValue("(Nii)", mangled, lineno, col_offset);
     if (!data)
         return 0;
     res = PyList_Append(st->st_cur->ste_directives, data);
@@ -1280,7 +1289,7 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
             }
             if (!symtable_add_def(st, name, DEF_GLOBAL))
                 VISIT_QUIT(st, 0);
-            if (!symtable_record_directive(st, name, s))
+            if (!symtable_record_directive(st, name, s->lineno, s->col_offset))
                 VISIT_QUIT(st, 0);
         }
         break;
@@ -1312,7 +1321,7 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
             }
             if (!symtable_add_def(st, name, DEF_NONLOCAL))
                 VISIT_QUIT(st, 0);
-            if (!symtable_record_directive(st, name, s))
+            if (!symtable_record_directive(st, name, s->lineno, s->col_offset))
                 VISIT_QUIT(st, 0);
         }
         break;
@@ -1368,6 +1377,60 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
 }
 
 static int
+symtable_extend_namedexpr_scope(struct symtable *st, expr_ty e)
+{
+    assert(st->st_stack);
+
+    Py_ssize_t i, size;
+    struct _symtable_entry *ste;
+    size = PyList_GET_SIZE(st->st_stack);
+    assert(size);
+
+    /* Iterate over the stack in reverse and add to the nearest adequate scope */
+    for (i = size - 1; i >= 0; i--) {
+        ste = (struct _symtable_entry *) PyList_GET_ITEM(st->st_stack, i);
+
+        /* If our current entry is a comprehension, skip it */
+        if (ste->ste_comprehension) {
+            continue;
+        }
+
+        /* If we find a FunctionBlock entry, add as NONLOCAL/LOCAL */
+        if (ste->ste_type == FunctionBlock) {
+            if (!symtable_add_def(st, e->v.Name.id, DEF_NONLOCAL))
+                VISIT_QUIT(st, 0);
+            if (!symtable_record_directive(st, e->v.Name.id, e->lineno, e->col_offset))
+                VISIT_QUIT(st, 0);
+
+            return symtable_add_def_helper(st, e->v.Name.id, DEF_LOCAL, ste);
+        }
+        /* If we find a ModuleBlock entry, add as GLOBAL */
+        if (ste->ste_type == ModuleBlock) {
+            if (!symtable_add_def(st, e->v.Name.id, DEF_GLOBAL))
+                VISIT_QUIT(st, 0);
+            if (!symtable_record_directive(st, e->v.Name.id, e->lineno, e->col_offset))
+                VISIT_QUIT(st, 0);
+
+            return symtable_add_def_helper(st, e->v.Name.id, DEF_GLOBAL, ste);
+        }
+        /* Disallow usage in ClassBlock */
+        if (ste->ste_type == ClassBlock) {
+            PyErr_Format(PyExc_TargetScopeError, NAMED_EXPR_COMP_IN_CLASS, e->v.Name.id);
+            PyErr_SyntaxLocationObject(st->st_filename,
+                                        e->lineno,
+                                        e->col_offset);
+            VISIT_QUIT(st, 0);
+        }
+    }
+
+    /* We should always find either a FunctionBlock, ModuleBlock or ClassBlock
+       and should never fall to this case
+    */
+    assert(0);
+    return 0;
+}
+
+static int
 symtable_visit_expr(struct symtable *st, expr_ty e)
 {
     if (++st->recursion_depth > st->recursion_limit) {
@@ -1376,6 +1439,10 @@ symtable_visit_expr(struct symtable *st, expr_ty e)
         VISIT_QUIT(st, 0);
     }
     switch (e->kind) {
+    case NamedExpr_kind:
+        VISIT(st, expr, e->v.NamedExpr.value);
+        VISIT(st, expr, e->v.NamedExpr.target);
+        break;
     case BoolOp_kind:
         VISIT_SEQ(st, expr, e->v.BoolOp.values);
         break;
@@ -1476,6 +1543,11 @@ symtable_visit_expr(struct symtable *st, expr_ty e)
         VISIT(st, expr, e->v.Starred.value);
         break;
     case Name_kind:
+        /* Special-case: named expr */
+        if (e->v.Name.ctx == NamedStore && st->st_cur->ste_comprehension) {
+            if(!symtable_extend_namedexpr_scope(st, e))
+                VISIT_QUIT(st, 0);
+        }
         if (!symtable_add_def(st, e->v.Name.id,
                               e->v.Name.ctx == Load ? USE : DEF_LOCAL))
             VISIT_QUIT(st, 0);
@@ -1713,6 +1785,8 @@ symtable_handle_comprehension(struct symtable *st, expr_ty e,
     if (outermost->is_async) {
         st->st_cur->ste_coroutine = 1;
     }
+    st->st_cur->ste_comprehension = 1;
+
     /* Outermost iter is received as an argument */
     if (!symtable_implicit_arg(st, 0)) {
         symtable_exit_block(st, (void *)e);