;;;;;; 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") }

(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) (index_expression_list) (exception_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

{ 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") }

(template_string string_content: (_) @part)

{ let @part.node = (ast-node @part "TemplateStringPart") }

; 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

}

}

(template_string) @tstring

{ let @tstring.node = (ast-node @tstring "TemplateString") }

(concatenated_template_string) @tstrings

{ let @tstrings.node = (ast-node @tstrings "JoinedTemplateString") }

(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)* . [(expression) (list_splat)] @start [(for_in_clause) (if_clause)] @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)* . [(expression) (list_splat)] @start (for_in_clause) @child [(for_in_clause) (if_clause)] @end . (comment)* . ")" .) @_genexpr

{

attr (@child.node) _location_start = (location-start @start)

attr (@child.node) _location_end = (location-end @end)

}

(generator_expression . "(" . (comment)* . [(expression) (list_splat)] @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) @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)* . [(expression) (list_splat)] @start [(for_in_clause) (if_clause)] @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

let @if.first_if = @if.node

}