bpo-39235: Fix end location for genexp in call args (GH-17925)

The fix changes copy_location() to require an extra node from which to extract the end location, and fixing all 5 call sites.


https://bugs.python.org/issue39235

2 files changed, 10 insertions, 8 deletions

diff --git a/Misc/NEWS.d/next/Core and Builtins/2020-01-09-10-01-18.bpo-39235.RYwjoc.rst b/Misc/NEWS.d/next/Core and Builtins/2020-01-09-10-01-18.bpo-39235.RYwjoc.rst
new file mode 100644
index 0000000..5fb0d45
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2020-01-09-10-01-18.bpo-39235.RYwjoc.rst
@@ -0,0 +1,2 @@
+Fix AST end location for lone generator expression in function call, e.g.
+f(i for i in a).
diff --git a/Python/ast.c b/Python/ast.c
index d5113f8..9c48d71 100644
--- a/Python/ast.c
+++ b/Python/ast.c
@@ -1028,13 +1028,13 @@ forbidden_name(struct compiling *c, identifier name, const node *n,
 }
 
 static expr_ty
-copy_location(expr_ty e, const node *n)
+copy_location(expr_ty e, const node *n, const node *end)
 {
     if (e) {
         e->lineno = LINENO(n);
         e->col_offset = n->n_col_offset;
-        e->end_lineno = n->n_end_lineno;
-        e->end_col_offset = n->n_end_col_offset;
+        e->end_lineno = end->n_end_lineno;
+        e->end_col_offset = end->n_end_col_offset;
     }
     return e;
 }
@@ -2464,10 +2464,10 @@ ast_for_atom(struct compiling *c, const node *n)
         }
 
         if (TYPE(CHILD(ch, 1)) == comp_for) {
-            return copy_location(ast_for_genexp(c, ch), n);
+            return copy_location(ast_for_genexp(c, ch), n, n);
         }
         else {
-            return copy_location(ast_for_testlist(c, ch), n);
+            return copy_location(ast_for_testlist(c, ch), n, n);
         }
     case LSQB: /* list (or list comprehension) */
         ch = CHILD(n, 1);
@@ -2486,7 +2486,7 @@ ast_for_atom(struct compiling *c, const node *n)
                         n->n_end_lineno, n->n_end_col_offset, c->c_arena);
         }
         else {
-            return copy_location(ast_for_listcomp(c, ch), n);
+            return copy_location(ast_for_listcomp(c, ch), n, n);
         }
     case LBRACE: {
         /* dictorsetmaker: ( ((test ':' test | '**' test)
@@ -2527,7 +2527,7 @@ ast_for_atom(struct compiling *c, const node *n)
                 /* It's a dictionary display. */
                 res = ast_for_dictdisplay(c, ch);
             }
-            return copy_location(res, n);
+            return copy_location(res, n, n);
         }
     }
     default:
@@ -3146,7 +3146,7 @@ ast_for_call(struct compiling *c, const node *n, expr_ty func,
             }
             else if (TYPE(CHILD(ch, 1)) == comp_for) {
                 /* the lone generator expression */
-                e = copy_location(ast_for_genexp(c, ch), maybegenbeg);
+                e = copy_location(ast_for_genexp(c, ch), maybegenbeg, closepar);
                 if (!e)
                     return NULL;
                 asdl_seq_SET(args, nargs++, e);