Python: Allow comments in comprehensions

A somewhat complicated solution that necessitated adding a new custom
function to `tsg-python`. See the comments in `python.tsg` for why this
was necessary.

Author: tausbn

python/extractor/tests/parser/comprehensions.py

@@ -55,3 +55,13 @@
 t = tuple(x for y in z)
 
 [(  t,  ) for v in w]
+
+[# comment
+    a for b in c # comment
+    # comment
+] # comment
+
+[# comment
+    d for e in f if g # comment
+    # comment
+] # comment

python/extractor/tsg-python/python.tsg

@@ -1062,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")
 
@@ -1096,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 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
+    if (instance-of last "if_clause") {
+        edge last.first_if -> expr
+        attr (last.first_if -> expr) body = 0
 
-    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

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