Merge pull request #4488 from charliermarsh/charlie/gen-exp-arg

youknowone · web-flow · commit 61b48f108982 · 2023-02-11T05:45:25.000+09:00
Use entire range for generators-as-arguments
diff --git a/.cspell.json b/.cspell.json
@@ -30,6 +30,7 @@
     "bindgen",
     "cstring",
     "chrono",
+    "insta",
     "peekable",
     "lalrpop",
     "memmap",
diff --git a/compiler/parser/python.lalrpop b/compiler/parser/python.lalrpop
@@ -1274,11 +1274,11 @@ ArgumentList: ArgumentList = {
 };
 
 FunctionArgument: (Option<(ast::Location, ast::Location, Option<String>)>, ast::Expr) = {
-    <e:NamedExpressionTest> <c:CompFor?> => {
+    <location:@L> <e:NamedExpressionTest> <c:CompFor?> <end_location:@R> => {
         let expr = match c {
             Some(c) => ast::Expr {
-                location: e.location,
-                end_location: e.end_location,
+                location,
+                end_location: Some(end_location),
                 custom: (),
                 node: ast::ExprKind::GeneratorExp {
                     elt: Box::new(e),
diff --git a/compiler/parser/src/parser.rs b/compiler/parser/src/parser.rs
@@ -310,6 +310,19 @@ with (0 as a, 1 as b,): pass
         }
     }
 
+    #[test]
+    fn test_generator_expression_argument() {
+        let source = r#"' '.join(
+    sql
+    for sql in (
+        "LIMIT %d" % limit if limit else None,
+        ("OFFSET %d" % offset) if offset else None,
+    )
+)"#;
+        let parse_ast = parse_expression(source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
     #[test]
     fn test_dict_unpacking() {
         let parse_ast = parse_expression(r#"{"a": "b", **c, "d": "e"}"#, "<test>").unwrap();
diff --git a/compiler/parser/src/snapshots/rustpython_parser__parser__tests__generator_expression_argument.snap b/compiler/parser/src/snapshots/rustpython_parser__parser__tests__generator_expression_argument.snap
