python.tsg

;;;;;; Part 1: Definining ~all~ most of the nodes
; This section contains all of the "simple" definitions. All of the places where a single
; tree-sitter node corresponds to an AST node.

; Create the module node first, so it always appears first in the output.
(module) @mod
{ let @mod.node = (ast-node @mod "Module") }

(_) @anynode
{
    scan (node-type @anynode) {
        "^(ERROR|MISSING)$" {
            let @anynode.node = (ast-node @anynode "SyntaxErrorNode")
            attr (@anynode.node) source = (source-text @anynode)
        }
    }
}

(parenthesized_expression) @nd
{ let @nd.node = (ast-node @nd "Expr") }

(assignment !type) @assign
{ let @assign.node = (ast-node @assign "Assign") }

[ (expression_list) (tuple) (tuple_pattern) (pattern_list) ] @tuple
{ let @tuple.node = (ast-node @tuple "Tuple") }

(call) @call { let @call.node = (ast-node @call "Call") }

(for_statement) @for
{ let @for.node = (ast-node @for "For") }

[ (if_statement) (elif_clause) ] @if
{ let @if.node = (ast-node @if "If") }

(continue_statement) @continue
{ let @continue.node = (ast-node @continue "Continue") }

(break_statement) @break
{ let @break.node = (ast-node @break "Break") }

(pass_statement) @pass
{ let @pass.node = (ast-node @pass "Pass") }

(assert_statement) @assert
{ let @assert.node = (ast-node @assert "Assert") }

(assignment type: (_)) @assign
{ let @assign.node = (ast-node @assign "AnnAssign") }

(augmented_assignment) @assign
{ let @assign.node = (ast-node @assign "AugAssign") }

(delete_statement) @del
{ let @del.node = (ast-node @del "Delete") }

(global_statement) @global
{ let @global.node = (ast-node @global "Global") }

(nonlocal_statement) @nonlocal
{ let @nonlocal.node = (ast-node @nonlocal "Nonlocal") }

[(import_statement) (import_from_statement name: (_))] @import
{ let @import.node = (ast-node @import "Import") }

(import_from_statement (wildcard_import)) @importstar
{ let @importstar.node = (ast-node @importstar "ImportFrom") }

(raise_statement) @raise
{ let @raise.node = (ast-node @raise "Raise") }

(binary_operator) @binop
{ let @binop.node = (ast-node @binop "BinOp") }

(keyword_argument) @kwarg
{ let @kwarg.node = (ast-node @kwarg "keyword") }

[(function_definition) (class_definition) (decorated_definition)] @def
{ let @def.node = (ast-node @def "Assign") }

(decorator) @decorator
{ let @decorator.node = (ast-node @decorator "Call") }

(expression_statement) @stmt
{ let @stmt.node = (ast-node @stmt "Expr") }

[ (integer) (float) ] @num
{ let @num.node = (ast-node @num "Num") }

(identifier) @name
{ let @name.node = (ast-node @name "Name") }

(list) @list
{ let @list.node = (ast-node @list "List") }

[(list_splat) (list_splat_pattern)] @starred
{ let @starred.node = (ast-node @starred "Starred") }

(comment) @comment
{ let @comment.node = (ast-node @comment "Comment") }

[
    (future_import_statement name: (_) @alias)
    (import_from_statement name: (_) @alias)
    (import_statement name: (_) @alias)
]
{ let @alias.node = (ast-node @alias "alias") }

; A string _without_ interpolations is just a `Str`, _except_ if it's inside a string
; concatenation, in which case it's a `StringPart`.
(string !interpolation) @str
{
    var str_class = "Str"
    if (instance-of (get-parent @str) "concatenated_string") {
        set str_class = "StringPart"
    }
    let @str.node = (ast-node @str str_class)
}

(string interpolation: (_)) @fstring
{ let @fstring.node = (ast-node @fstring "JoinedStr") }


(string string_content: (_) @part)
{ let @part.node = (ast-node @part "StringPart") }

; A string concatenation that contains no interpolated expressions is just a `Str` (and its children
; will be `StringPart`s). A string concatenation that contains interpolated expressions is a
; `JoinedStr`, however.
(concatenated_string
    (string interpolation: (_))* @interpolations
) @string
{
    var string_class = "Str"
    ; Check if there are any interpolations in the string.
    ; We cannot use an optional match in the above query, since it could match several times,
    ; and subsequent definitions of `@string.node` would then fail.
    for _ in @interpolations {
        set string_class = "JoinedStr"
    }
    let @string.node = (ast-node @string string_class)
}


(string interpolation: (_)) @fstring
{
    if (not (instance-of (get-parent @fstring) "concatenated_string")) {
        attr (@fstring.node) _fixup = #true
    }
}

(pair) @kvpair
{ let @kvpair.node = (ast-node @kvpair "KeyValuePair") }

(dictionary) @dict
{ let @dict.node = (ast-node @dict "Dict") }

(dictionary_splat) @dictunpacking
{ let @dictunpacking.node = (ast-node @dictunpacking "DictUnpacking") }

(set) @set
{ let @set.node = (ast-node @set "Set") }

(boolean_operator) @boolop
{ let @boolop.node = (ast-node @boolop "BoolOp") }

(comparison_operator) @compop
{ let @compop.node = (ast-node @compop "Compare") }

[ (unary_operator) (not_operator) ] @unaryop
{ let @unaryop.node = (ast-node @unaryop "UnaryOp") }

(exec_statement) @exec
{ let @exec.node = (ast-node @exec "Exec") }

(print_statement) @print
{ let @print.node = (ast-node @print "Print") }

(return_statement) @return
{ let @return.node = (ast-node @return "Return") }

(yield . "from"? @from) @yield
{
    var yield_node = "Yield"
    if some @from {
        set yield_node = "YieldFrom"
    }
    let @yield.node = (ast-node @yield yield_node)
}

(ellipsis) @ellipsis
{ let @ellipsis.node = (ast-node @ellipsis "Ellipsis") }

(await) @await
{ let @await.node = (ast-node @await "Await") }

(try_statement) @try
{ let @try.node = (ast-node @try "Try") }

(except_clause) @except
{ let @except.node = (ast-node @except "ExceptStmt") }

(except_group_clause) @except
{ let @except.node = (ast-node @except "ExceptGroupStmt") }

(named_expression) @assignexpr
{ let @assignexpr.node = (ast-node @assignexpr "AssignExpr") }

(conditional_expression) @ifexp
{ let @ifexp.node = (ast-node @ifexp "IfExp") }

(subscript) @subscript
{ let @subscript.node = (ast-node @subscript "Subscript") }

(slice) @slice
{ let @slice.node = (ast-node @slice "Slice") }

(attribute) @attribute
{ let @attribute.node = (ast-node @attribute "Attribute") }

(while_statement) @while
{ let @while.node = (ast-node @while "While") }

(generator_expression) @generatorexp
{ let @generatorexp.node = (ast-node @generatorexp "GeneratorExp") }

(for_in_clause) @for
{ let @for.node = (ast-node @for "For") }

(if_clause) @if
{ let @if.node = (ast-node @if "If") }

(list_comprehension) @listcomp
{ let @listcomp.node = (ast-node @listcomp "ListComp") }

(set_comprehension) @setcomp
{ let @setcomp.node = (ast-node @setcomp "SetComp") }

(dictionary_comprehension) @dictcomp
{ let @dictcomp.node = (ast-node @dictcomp "DictComp") }

[ (with_statement) (with_item)] @with
{ let @with.node = (ast-node @with "With") }

(match_statement) @match
{ let @match.node = (ast-node @match "Match") }

; Do not create an AST node for 'cases', we just wire up the children instead.

(case_block) @case
{ let @case.node = (ast-node @case "Case") }

(match_as_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchAsPattern") }

(match_or_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchOrPattern") }

(match_literal_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchLiteralPattern") }

(match_capture_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchCapturePattern") }

(match_wildcard_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchWildcardPattern") }

(match_value_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchValuePattern") }

(match_group_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchGroupPattern") }

(match_sequence_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchSequencePattern") }

(match_star_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchStarPattern") }

(match_mapping_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchMappingPattern") }

(match_double_star_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchDoubleStarPattern") }

(match_key_value_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchKeyValuePattern") }

(match_class_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchClassPattern") }

; Do not create AST nodes for 'only_positionals', 'only_keywords',
; 'partly_positionals', and 'partly_keywords'. We just wire up the children instead.

(match_keyword_pattern) @pattern
{ let @pattern.node = (ast-node @pattern "MatchKeywordPattern") }

(guard) @guard
{ let @guard.node = (ast-node @guard "Guard") }

[(parameters) (lambda_parameters)] @params
{ let @params.node = (ast-node @params "arguments") }

[(false) (true) (none)] @const
{ let @const.node = (ast-node @const "Name") }

(lambda) @lambda
{ let @lambda.node = (ast-node @lambda "Lambda") }

(future_import_statement) @import
{ let @import.node = (ast-node @import "Import") }

(typevar_parameter) @typevar
{ let @typevar.node = (ast-node @typevar "TypeVar") }

(typevartuple_parameter) @typevartuple
{ let @typevartuple.node = (ast-node @typevartuple "TypeVarTuple") }

(paramspec_parameter) @paramspec
{ let @paramspec.node = (ast-node @paramspec "ParamSpec") }

(type_alias_statement) @typealias
{ let @typealias.node = (ast-node @typealias "TypeAlias") }

;;;;;; End of part 1.

;;;;;; Part 2: The awkward bunch.

;;;;;; Workarounds for node locations
; These are (hopefully temporary) workarounds for the nodes for which the default start and end does
; not agree with what our internal AST provides.
; Once the new parser is in place, we can consider getting rid of these workarounds.


;;; If
; End position is set to the end of the `:` after the condition.
[
    (if_statement
        condition: (_)
        .
        ":" @colon) @if
    (elif_clause
        condition: (_)
        .
        ":" @colon) @if
]
{
    attr (@if.node) _location_end = (location-end @colon)
}

;;; For
; Same as with `if`, we must include the `:` in the position.
(for_statement
    right: (_)
    .
    ":" @colon
) @for
{
    attr (@for.node) _location_end = (location-end @colon)
}

;;; While
; Same as with `if`, we must include the `:` in the position.
(while_statement
    condition: (_)
    .
    ":" @colon
) @while
{
    attr (@while.node) _location_end = (location-end @colon)
}

;;; Tuples
; In the Python AST tuple start and end positions are set to the start and end of the first and last
; elements. In `tree-sitter-python`, the parentheses are included.
[
    (tuple . (comment)* . element: (_) @first)
    (tuple_pattern . (comment)* . element: (_) @first)
] @tuple
{
    attr (@tuple.node) _location_start = (location-start @first)
}

[
    (tuple !trailing_comma element: (_) @last . (comment)* . ")" .)
    (tuple trailing_comma: _ @last)
    (tuple_pattern element: (_) @last .)
] @tuple
{

    attr (@tuple.node) _location_end = (location-end @last)
}

;;; Try

(try_statement ":" @colon) @try
{ attr (@try.node) _location_end = (location-end @colon) }

(except_clause ":" @colon) @except
{ attr (@except.node) _location_end = (location-end @colon) }

;;; GeneratorExp

(generator_expression . "(" . (comment)* . (_) @start (_) @end . (comment)* . ")" .) @generatorexp
{
    attr (@generatorexp.node) _location_start = (location-start @start)
    attr (@generatorexp.node) _location_end = (location-end @end)
}

(if_clause (expression) @expr) @if
{
    attr (@if.node) _location_start = (location-start @expr)
    attr (@if.node) _location_end = (location-end @expr)
}

(generator_expression . "(" . (comment)* . (_) @start (for_in_clause) @child (_) @end . (comment)* . ")" .) @genexpr
{
    attr (@child.node) _location_start = (location-start @start)
    attr (@child.node) _location_end = (location-end @end)
}

(generator_expression . "(" . (comment)* . (_) @start (for_in_clause) @end . (comment)* . ")" .) @genexpr
{
    attr (@end.node) _location_start = (location-start @start)
    attr (@end.node) _location_end = (location-end @end)
}

(list_comprehension (for_in_clause) @child) @genexpr
{
    attr (@child.node) _location_start = (location-start @genexpr)
    attr (@child.node) _location_end = (location-end @genexpr)
}

(set_comprehension  (for_in_clause) @child) @genexpr
{
    attr (@child.node) _location_start = (location-start @genexpr)
    attr (@child.node) _location_end = (location-end @genexpr)
}

(dictionary_comprehension (for_in_clause) @child) @genexpr
{
    attr (@child.node) _location_start = (location-start @genexpr)
    attr (@child.node) _location_end = (location-end @genexpr)
}

;;; With

(with_statement
    "with" @start
    (with_clause . (with_item) @first)
    ":" @end
)
{
    attr (@first.node) _location_start = (location-start @start)
    attr (@first.node) _location_end = (location-end @end)
}

;;;;;; End of workarounds

;;;;;; End of part 2

;;;;;; Part 3: All of the simple nodes.

;;;;;; Module

; Nodes with a `body` field containing statements.
(module (_) @stmt) @parent
{
    edge @parent.node -> @stmt.node
    attr (@parent.node -> @stmt.node) body = (named-child-index @stmt)
}

;;;;;; Comments

(comment) @comment
{
    attr (@comment.node) text = (source-text @comment)
}

;;;;;; Expressions

(parenthesized_expression
    inner: (_) @inner
) @outer
{
    attr (@outer.node) _skip_to = @inner.node
    attr (@inner.node) parenthesised = #true
}

(keyword_argument
    name: (_) @name
    value: (_) @value
) @kwarg
{
    attr (@kwarg.node) arg = (source-text @name)
    attr (@kwarg.node) value = @value.node
}

;;;;;; Num

[ (integer) (float) ] @num
{
    ; As we must support a large variety of number literals, we simply forward the source string
    ; representation to the Python AST reconstruction.
    let source = (source-text @num)
    attr (@num.node) n = source
    attr (@num.node) text = source
}

;;;;;; End of Num

;;;;;; Delete

(delete_statement
    target: (expression_list
        element: (_) @target
    )
) @del
{
    edge @del.node -> @target.node
    attr (@del.node -> @target.node) targets = (named-child-index @target)
}

(delete_statement target: (_) @target) @del
{
    attr (@target.node) ctx = "del"
}

(delete_statement [(identifier) (subscript) (attribute)] @id) @del
{
    edge @del.node -> @id.node
    attr (@del.node -> @id.node) targets = 0
}

;;;;;; Name

[(identifier) (false) (true) (none)] @id
{
    attr (@id.node) variable = (source-text @id)
}

;;;;;; End of Name

;;;;;; Arguments
[
    (keyword_argument value: (_) @id)
    (argument_list element: (_) @id)
]
{
    attr (@id.node) ctx = "load"
}

[
    (keyword_argument name: (_) @id)
]
{
    attr (@id.node) ctx = "store"
}

;;;;;; End of Arguments

;;;;;; BinOp

(binary_operator
    left: (_) @left
    operator: _ @op
    right: (_) @right
) @bin
{
    attr (@bin.node) left = @left.node
    attr (@bin.node) right = @right.node
    attr (@bin.node) op = (source-text @op)
    attr (@left.node) ctx = "load"
    attr (@right.node) ctx = "load"
}

;;;;;; End of BinOp

;;;;;; If

; If.test
[
    (if_statement
        condition: (_) @test) @if
    (elif_clause
        condition: (_) @test) @if
]
{
    attr (@if.node) test = @test.node
    attr (@test.node) ctx = "load"
}

; If.orelse - first `elif` clause
(if_statement
    consequence: (_)
    . (comment)* .
    (elif_clause) @elif
) @if
{
    edge @if.node -> @elif.node
    attr (@if.node -> @elif.node) orelse = 0
}

; If.orelse - link up adjacent `elif` clauses
(
    (elif_clause) @elif1
    . (comment)* .
    (elif_clause) @elif2
)
{
    edge @elif1.node -> @elif2.node
    attr (@elif1.node -> @elif2.node) orelse = 0
}

; If.orelse - match outer `else` up with last `elif` clause (i.e. innermost `if`)
(if_statement
    (elif_clause) @elif . (comment)* .
    alternative: (else_clause body: (block (_) @orelse))
)
{
    edge @elif.node -> @orelse.node
    attr (@elif.node -> @orelse.node) orelse = (named-child-index @orelse)
}

; If.orelse - when there are no `elif` clauses.
(if_statement
    consequence: (_)
    . (comment)* .
    alternative: (else_clause body: (block (_) @orelse))
) @if
{
    edge @if.node -> @orelse.node
    attr (@if.node -> @orelse.node) orelse = (named-child-index @orelse)
}

; If.body

[
    (if_statement
        consequence: (block (_) @stmt)) @parent
    (elif_clause
        consequence: (block (_) @stmt)) @parent
]
{
    edge @parent.node -> @stmt.node
    attr (@parent.node -> @stmt.node) body = (named-child-index @stmt)
}

;;;;;; end of If

;;;;;; For statements

(for_statement
    left: (_) @left
    right: (_) @right
) @for
{
    attr (@for.node) target = @left.node
    attr (@left.node) ctx = "store"
    attr (@for.node) iter = @right.node
    attr (@right.node) ctx = "load"
}

(for_statement
    body: (block (_) @body)
) @for
{
    edge @for.node -> @body.node
    attr (@for.node -> @body.node) body = (named-child-index @body)
}

(for_statement
    alternative: (else_clause body: (block (_) @orelse))
) @for
{
    edge @for.node -> @orelse.node
    attr (@for.node -> @orelse.node) orelse = (named-child-index @orelse)
}

(for_statement "async" "for" @for_keyword) @for
{
    attr (@for.node) is_async = #true
    attr (@for.node) _location_start = (location-start @for_keyword)
}

;;;;;; end of For

;;;;;; Call expressions (`a(b, c, *d, **e)`)

(call function: (_) @func) @call
{
    attr (@call.node) func = @func.node
    attr (@func.node) ctx = "load"
}

; Handle non-keyword arguments
(call arguments: (argument_list element: (_) @arg)) @call
{
    if (not (or
            (instance-of @arg "keyword_argument")
            (instance-of @arg "dictionary_splat"))) {
        edge @call.node -> @arg.node
        attr (@call.node -> @arg.node) positional_args = (named-child-index @arg)
    }
}

(call arguments: (argument_list element: (keyword_argument) @arg)) @call
{
    edge @call.node -> @arg.node
    attr (@call.node -> @arg.node) named_args = (named-child-index @arg)
}

(call arguments: (argument_list element: (dictionary_splat) @arg)) @call
{
    edge @call.node -> @arg.node
    attr (@call.node -> @arg.node) named_args = (named-child-index @arg)
}

(call arguments: (generator_expression) @gen) @call
{
    edge @call.node -> @gen.node
    attr (@call.node -> @gen.node) positional_args = 0
}


;;;;;; end of Call (`a(b, c, *d, **e)`)


;;;;;; End of part 3

;;;;;; Part 4: All of the complicated bits (e.g. nodes that need additional synthesis)


;;;;;; ListComp (`[a for b in c if d]`)
; See GeneratorExp for details.

(list_comprehension) @genexpr
{
    ; Synthesize the `genexpr` function
    let @genexpr.fun = (ast-node @genexpr "Function")
    attr (@genexpr.node) function = @genexpr.fun
    attr (@genexpr.fun) name = "listcomp"

    ; Synthesize the `.0` parameter
    let @genexpr.arg = (ast-node @genexpr "Name")
    attr (@genexpr.arg) variable = ".0"
    attr (@genexpr.arg) ctx = "param"

    edge @genexpr.fun -> @genexpr.arg
    attr (@genexpr.fun -> @genexpr.arg) args = 0
    attr (@genexpr.fun) kwonlyargs = #null
    attr (@genexpr.fun) kwarg = #null

    ; Synthesize the use of `.0` in the outermost `for`. This has a different context than the parameter
    ; ("param" vs. "load") hence we must create another node.
    let @genexpr.arg_use = (ast-node @genexpr "Name")
    attr (@genexpr.arg_use) variable = ".0"
    attr (@genexpr.arg_use) ctx = "load"
}

;;;;;; End of ListComp (`[a for b in c if d]`)

;;;;;; SetComp (`{a for b in c if d}`)
; See GeneratorExp for details.

(set_comprehension) @genexpr
{
    ; Synthesize the `genexpr` function
    let @genexpr.fun = (ast-node @genexpr "Function")
    attr (@genexpr.node) function = @genexpr.fun
    attr (@genexpr.fun) name = "setcomp"

    ; Synthesize the `.0` parameter
    let @genexpr.arg = (ast-node @genexpr "Name")
    attr (@genexpr.arg) variable = ".0"
    attr (@genexpr.arg) ctx = "param"

    edge @genexpr.fun -> @genexpr.arg
    attr (@genexpr.fun -> @genexpr.arg) args = 0
    attr (@genexpr.fun) kwonlyargs = #null
    attr (@genexpr.fun) kwarg = #null

    ; Synthesize the use of `.0` in the outermost `for`. This has a different context than the parameter
    ; ("param" vs. "load") hence we must create another node.
    let @genexpr.arg_use = (ast-node @genexpr "Name")
    attr (@genexpr.arg_use) variable = ".0"
    attr (@genexpr.arg_use) ctx = "load"
}


;;;;;; End of SetComp (`{a for b in c if d}`)

;;;;;; DictComp (`{a: b for c in d if e}`)
; See GeneratorExp for details.

(dictionary_comprehension
    body: (pair
        key: (_) @key
        value: (_) @value
    )
) @genexpr
{
    ; Synthesize the `genexpr` function
    let @genexpr.fun = (ast-node @genexpr "Function")
    attr (@genexpr.node) function = @genexpr.fun
    attr (@genexpr.fun) name = "dictcomp"

    ; Synthesize the `.0` parameter
    let @genexpr.arg = (ast-node @genexpr "Name")
    attr (@genexpr.arg) variable = ".0"
    attr (@genexpr.arg) ctx = "param"

    edge @genexpr.fun -> @genexpr.arg
    attr (@genexpr.fun -> @genexpr.arg) args = 0
    attr (@genexpr.fun) kwonlyargs = #null
    attr (@genexpr.fun) kwarg = #null

    ; Synthesize the use of `.0` in the innermost `yield`. This has a different context than the parameter
    ; ("param" vs. "load") hence we must create another node.
    let @genexpr.arg_use = (ast-node @genexpr "Name")
    attr (@genexpr.arg_use) variable = ".0"
    attr (@genexpr.arg_use) ctx = "load"
}

;;;;;; End of DictComp (`{a: b for c in d if e}`)

;;;;;; GeneratorExp (`(a for b in c if d)`)
; The big one. This one will require quite a bit of setup.
;
; First of all, we need to explain what the old parser does to generator expressions.
;
; The following generator expression
;
; (a
;     for b in c
;         if d
;         if e
;     for f in g
;         if h
;         if i
; )
;
; becomes
;
; def genexpr(.0):
;     for b in .0:
;         if e:
;            if d:
;                for f in g:
;                   if i:
;                       if h:
;                           yield a
;
; where `.0` is a (very oddly named) variable.
;
; Note in particular the reversing of the `if`s, the way `c` is replaced with `.0`, and the way
; `a` is used in the innermost `yield`.

; First of all, we need to set up the generated function and its parameter. These both copy the location
; information for the entire generator expression (yes, it is a wide parameter!) and so we must recreate the logic for
; setting this location information correctly.

(generator_expression . "(" . (comment)* . (_) @start  (_) @end . (comment)* . ")" .) @genexpr
{
    ; Synthesize the `genexpr` function
    let @genexpr.fun = (ast-node @genexpr "Function")
    attr (@genexpr.fun) _location_start = (location-start @start)
    attr (@genexpr.fun) _location_end = (location-end @end)
    attr (@genexpr.node) function = @genexpr.fun
    attr (@genexpr.fun) name = "genexpr"

    ; Synthesize the `.0` parameter
    let @genexpr.arg = (ast-node @genexpr "Name")
    attr (@genexpr.arg) _location_start = (location-start @start)
    attr (@genexpr.arg) _location_end = (location-end @end)
    attr (@genexpr.arg) variable = ".0"
    attr (@genexpr.arg) ctx = "param"

    edge @genexpr.fun -> @genexpr.arg
    attr (@genexpr.fun -> @genexpr.arg) args = 0
    attr (@genexpr.fun) kwonlyargs = #null
    attr (@genexpr.fun) kwarg = #null

    ; Default to true, but we'll set it to false if we're inside a call
    var genexpr_parenthesised = #true

    if (instance-of (get-parent @genexpr) "call") {
        set genexpr_parenthesised = #null
    }
    attr (@genexpr.node) parenthesised = genexpr_parenthesised

    ; Synthesize the use of `.0` in the outermost `for`. This has a different context than the parameter
    ; ("param" vs. "load") hence we must create another node.
    let @genexpr.arg_use = (ast-node @genexpr "Name")
    attr (@genexpr.arg_use) _location_start = (location-start @start)
    attr (@genexpr.arg_use) _location_end = (location-end @end)
    attr (@genexpr.arg_use) variable = ".0"
    attr (@genexpr.arg_use) ctx = "load"
}

; Link up the outermost `for`
[
    (generator_expression
        body: (_) . (comment)* .
        (for_in_clause
            left: (_) @target
            right: (_) @iterable
        ) @forin
    ) @genexpr
    (list_comprehension
        body: (_) . (comment)* .
        (for_in_clause
            left: (_) @target
            right: (_) @iterable
        ) @forin
    ) @genexpr
    (set_comprehension
        body: (_) . (comment)* .
        (for_in_clause
            left: (_) @target
            right: (_) @iterable
        ) @forin
    ) @genexpr
    (dictionary_comprehension
        body: (_) . (comment)* .
        (for_in_clause
            left: (_) @target
            right: (_) @iterable
        ) @forin
    ) @genexpr
]
{
    attr (@genexpr.node) iterable = @iterable.node
    attr (@iterable.node) ctx = "load"
    edge @genexpr.fun -> @forin.node
    attr (@genexpr.fun -> @forin.node) body = 0

    attr (@forin.node) target = @target.node
    attr (@target.node) ctx = "store"
    attr (@forin.node) iter = @genexpr.arg_use
}

; Set up all subsequent `for ... in ...`
[
    (generator_expression
        body: (_)
        [(for_in_clause) (if_clause)]
        (for_in_clause left: (_) @target right: (_) @iter) @forin
    )
    (list_comprehension
        body: (_)
        [(for_in_clause) (if_clause)]
        (for_in_clause left: (_) @target right: (_) @iter) @forin
    )
    (set_comprehension
        body: (_)
        [(for_in_clause) (if_clause)]
        (for_in_clause left: (_) @target right: (_) @iter) @forin
    )
    (dictionary_comprehension
        body: (_)
        [(for_in_clause) (if_clause)]
        (for_in_clause left: (_) @target right: (_) @iter) @forin
    )
]
{
    attr (@forin.node) target = @target.node
    attr (@target.node) ctx = "store"
    attr (@forin.node) iter = @iter.node
    attr (@iter.node) ctx = "load"
}

; Set up each `if ...`
(if_clause (expression) @test) @if
{
    attr (@if.node) test = @test.node
    attr (@test.node) ctx = "load"
}

; Link adjacent `for` clauses together
(_
    (for_in_clause) @forin1
    . (comment)* .
    (for_in_clause) @forin2
)
{
    edge @forin1.node -> @forin2.node
    attr (@forin1.node -> @forin2.node) body = 0
}

; For the first `if` clause after a `for` clause, record both the `for` and `if` clauses in variables that we
; will propagate along. That way, when we get to the last `if` clause, we can link it up with the `for`
; clause, and we can link up the _first_ `if` clause with whatever follows the last `if` clause.
(_
    (for_in_clause) @forin
    . (comment)* .
    (if_clause) @if
)
{
    let @if.for = @forin.node