Merge pull request #17822 from github/tausbn/python-more-parser-fixes

Python: A few more parser fixes

python/extractor/tsg-python/python.tsg

@@ -25,6 +25,9 @@
 [ (expression_list) (tuple) (tuple_pattern) (pattern_list) ] @tuple
 { let @tuple.node = (ast-node @tuple "Tuple") }
 
+(list_pattern) @list
+{ let @list.node = (ast-node @list "List") }
+
 (call) @call { let @call.node = (ast-node @call "Call") }
 
 (for_statement) @for
@@ -1059,30 +1062,38 @@
     let @genexpr.result = tuple
 }
 
-; For the final `if` clause, we need to hook it up with the `yield` expression and with its associated `for` clause.
+; For the final clause, we need to hook it up with the rest of the expression.
+; If it's an `if` clause, we need to hook it up with the `yield` expression and with its associated
+; `for` clause.
+; If it's a `for` clause, we only need to create and hook it up with the `yield` expression.
+;
+; It would be tempting to use anchors here, but they just don't work. In particular, an anchor of
+; the form `. (comment)* . )` (which would be needed in order to handle the case where there are
+; comments after the last clause) cause the `tree-sitter` query engine to match _all_ clauses, not
+; just the last one.
+; Instead, we gather up all clauses in a list (these will be in the order they appear in the source
+; code), and extract the last element using a custom Rust function.
 [
     (generator_expression
         body: (_) @body
-        (if_clause) @last
-        .
+        [(if_clause) (for_in_clause)]+ @last_candidates
     ) @genexpr
     (list_comprehension
         body: (_) @body
-        (if_clause) @last
-        .
+        [(if_clause) (for_in_clause)]+ @last_candidates
     ) @genexpr
     (set_comprehension
         body: (_) @body
-        (if_clause) @last
-        .
+        [(if_clause) (for_in_clause)]+ @last_candidates
     ) @genexpr
     (dictionary_comprehension
         body: (_) @body
-        (if_clause) @last
-        .
+        [(if_clause) (for_in_clause)]+ @last_candidates
     ) @genexpr
 ]
 {
+    let last = (get-last-element @last_candidates)
+
     let expr = (ast-node @body "Expr")
     let yield = (ast-node @body "Yield")
 
@@ -1093,50 +1104,19 @@
 
     attr (yield) value = @genexpr.result
     attr (@body.node) ctx = "load"
-    edge @last.first_if -> expr
-    attr (@last.first_if -> expr) body = 0
 
-    ; Hook up this `if` clause with its `for` clause
-    edge @last.for -> @last.node
-    attr (@last.for -> @last.node) body = 0
-}
+    if (instance-of last "if_clause") {
+        edge last.first_if -> expr
+        attr (last.first_if -> expr) body = 0
 
-; If the last clause is a `for`, we only have to create and hook up the `yield` expression.
-[
-    (generator_expression
-        body: (_) @body
-        (for_in_clause) @last
-        .
-    ) @genexpr
-    (list_comprehension
-        body: (_) @body
-        (for_in_clause) @last
-        .
-    ) @genexpr
-    (set_comprehension
-        body: (_) @body
-        (for_in_clause) @last
-        .
-    ) @genexpr
-    (dictionary_comprehension
-        body: (_) @body
-        (for_in_clause) @last
-        .
-    ) @genexpr
-]
-{
-    let expr = (ast-node @body "Expr")
-    let yield = (ast-node @body "Yield")
-
-    let @genexpr.expr = expr
-    let @genexpr.yield = yield
-
-    attr (expr) value = yield
-
-    attr (yield) value = @genexpr.result
-    attr (@body.node) ctx = "load"
-    edge @last.node -> expr
-    attr (@last.node -> expr) body = 0
+        ; Hook up this `if` clause with its `for` clause
+        edge last.for -> last.node
+        attr (last.for -> last.node) body = 0
+    } else {
+        ; If the last clause is a `for`, we only have to create and hook up the `yield` expression.
+        edge last.node -> expr
+        attr (last.node -> expr) body = 0
+    }
 }
 
 ; For whatever reason, we do not consider parentheses around the yielded expression if they are present, so
@@ -3180,11 +3160,11 @@
             (typed_parameter
                 (identifier) @name
                 .
-                type: (type (expression) @type)
+                type: (type (_) @type)
             )
             (typed_default_parameter
                 name: (_) @name
-                type: (type (expression) @type)
+                type: (type (_) @type)
                 value: (_) @value
             )
         ] @param
@@ -3239,7 +3219,7 @@
             (list_splat_pattern vararg: (_) @name) @starred
             (typed_parameter
                 (list_splat_pattern vararg: (_) @name) @starred
-                type: (type (expression) @type)
+                type: (type (_) @type)
             )
         ]
     ) @params
@@ -3256,7 +3236,7 @@
 
 ; Return type
 (function_definition
-    return_type: (type (expression) @type)
+    return_type: (type (_) @type)
 ) @funcdef
 {
     attr (@funcdef.funcexpr) returns = @type.node
@@ -3270,7 +3250,7 @@
             (dictionary_splat_pattern kwarg: (identifier) @name)
             (typed_parameter
                 (dictionary_splat_pattern kwarg: (identifier) @name)
-                type: (type (expression) @type)
+                type: (type (_) @type)
             )
         ]
     ) @params
@@ -3447,6 +3427,9 @@
     ; Left hand side of an assignment such as `foo, bar = ...`
     (pattern_list element: (_) @elt) @parent
 
+    ; Left hand side of an assignment such as `[foo, bar] = ...`
+    (list_pattern element: (_) @elt) @parent
+
     ; An unadorned tuple (such as in `x = y, z`)
     (expression_list element: (_) @elt) @parent
 
@@ -3483,6 +3466,7 @@
     (tuple element: (_) @elt)
     (tuple_pattern element: (_) @elt)
     (pattern_list element: (_) @elt)
+    (list_pattern element: (_) @elt)
     (expression_list element: (_) @elt)
     (parenthesized_expression inner: (_) @elt)
     (set element: (_) @elt)

python/extractor/tsg-python/src/main.rs

@@ -463,6 +463,22 @@ pub mod extra_functions {
             Ok(Value::Integer(left % right))
         }
     }
+
+    pub struct GetLastElement;
+
+    impl Function for GetLastElement {
+        fn call(
+            &self,
+            _graph: &mut Graph,
+            _source: &str,
+            parameters: &mut dyn Parameters,
+        ) -> Result<Value, ExecutionError> {
+            let list = parameters.param()?.into_list()?;
+            parameters.finish()?;
+            let last = list.last().unwrap_or(&Value::Null).clone();
+            Ok(last)
+        }
+    }
 }
 
 fn main() -> Result<()> {
@@ -562,6 +578,12 @@ fn main() -> Result<()> {
     );
 
     functions.add(Identifier::from("mod"), extra_functions::Modulo);
+
+    functions.add(
+        Identifier::from("get-last-element"),
+        extra_functions::GetLastElement,
+    );
+
     let globals = Variables::new();
     let mut config = ExecutionConfig::new(&mut functions, &globals).lazy(false);
     let graph = file

python/extractor/tsg-python/tsp/grammar.js

@@ -309,7 +309,8 @@ module.exports = grammar({
     ),
 
     except_group_clause: $ => seq(
-      'except*',
+      'except',
+      '*',
       seq(
         field('type', $.expression),
         optional(seq(
@@ -963,7 +964,7 @@ module.exports = grammar({
       field('type', $.type)
     )),
 
-    type: $ => $.expression,
+    type: $ => choice($.list_splat, $.expression),
 
     keyword_argument: $ => seq(
       field('name', choice($.identifier, $.keyword_identifier)),

python/extractor/tsg-python/tsp/src/grammar.json

@@ -1169,7 +1169,11 @@
       "members": [
         {
           "type": "STRING",
-          "value": "except*"
+          "value": "except"
+        },
+        {
+          "type": "STRING",
+          "value": "*"
         },
         {
           "type": "SEQ",
@@ -5289,8 +5293,17 @@
       }
     },
     "type": {
-      "type": "SYMBOL",
-      "name": "expression"
+      "type": "CHOICE",
+      "members": [
+        {
+          "type": "SYMBOL",
+          "name": "list_splat"
+        },
+        {
+          "type": "SYMBOL",
+          "name": "expression"
+        }
+      ]
     },
     "keyword_argument": {
       "type": "SEQ",

python/extractor/tsg-python/tsp/src/node-types.json

@@ -3347,6 +3347,10 @@
         {
           "type": "expression",
           "named": true
+        },
+        {
+          "type": "list_splat",
+          "named": true
         }
       ]
     }
@@ -3960,10 +3964,6 @@
     "type": "except",
     "named": false
   },
-  {
-    "type": "except*",
-    "named": false
-  },
   {
     "type": "exec",
     "named": false