;;;;;; 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) (index_expression_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") }
; 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)* . (expression) @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) @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) @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)* . (expression) @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
}
; Link up adjacent `if` clauses (note the reversed order!) and propagate the `for` and `first_if` values.
(_
(if_clause) @if1
. (comment)* .
(if_clause) @if2
)
{
edge @if2.node -> @if1.node
attr (@if2.node -> @if1.node) body = 0
let @if2.for = @if1.for
let @if2.first_if = @if1.first_if
}
; After the last `if` in a chain, we hook it up as the body of its associated `for`, and hook up the _first_
; `if` as the one that has the following `for` as its body.
; The case where there is no `for` following the last `if` is handled later.
(_
(if_clause) @if
. (comment)* .
(for_in_clause) @forin
)
{
edge @if.for -> @if.node
attr (@if.for -> @if.node) body = 0
edge @if.first_if -> @forin.node
attr (@if.first_if -> @forin.node) body = 0
}
; For everything except dictionary comprehensions, the innermost expression is just the `body` of the
; comprehension.
[
(generator_expression body: (_) @body) @genexpr
(list_comprehension body: (_) @body) @genexpr
(set_comprehension body: (_) @body) @genexpr
]
{
let @genexpr.result = @body.node
}
; For dict comprehensions, we build an explicit tuple using the key and value pair.
(dictionary_comprehension
body: (pair
key: (_) @key
value: (_) @value
) @body
) @genexpr
{
let tuple = (ast-node @body "Tuple")
edge tuple -> @key.node
attr (tuple -> @key.node) elts = 1
edge tuple -> @value.node
attr (tuple -> @value.node) elts = 0
; TODO verify that it is correct to use a `(value, key)` tuple, and not a `(key, value)` tuple above.
; That is what the current parser does...
attr (tuple) ctx = "load"
let @genexpr.result = tuple
}
; 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) (for_in_clause)]+ @last_candidates
) @genexpr
(list_comprehension
body: (_) @body
[(if_clause) (for_in_clause)]+ @last_candidates
) @genexpr
(set_comprehension
body: (_) @body
[(if_clause) (for_in_clause)]+ @last_candidates
) @genexpr
(dictionary_comprehension
body: (_) @body
[(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")
let @genexpr.expr = expr
let @genexpr.yield = yield
attr (expr) value = yield
attr (yield) value = @genexpr.result
attr (@body.node) ctx = "load"
if (instance-of last "if_clause") {
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
} 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
; we must adapt the location accordingly.
[
(generator_expression
body: (_ . "(" . _ @first)
)
(list_comprehension
body: (_ . "(" . _ @first)
)
(set_comprehension
body: (_ . "(" . _ @first)
)
(dictionary_comprehension
body: (_ . "(" . _ @first)
)
] @genexpr
{
attr (@genexpr.expr) _location_start = (location-start @first)
attr (@genexpr.yield) _location_start = (location-start @first)
}
; Annoyingly, setting the end location of the synthesized `Expr` and `Yield` is a big mess,
; so we have to use mutable variables.
[
(generator_expression body: (_) @body)
(list_comprehension body: (_) @body)
(set_comprehension body: (_) @body)
(dictionary_comprehension body: (_) @body)
] @genexpr
{
var @genexpr.body_end = (location-end @body)
}
; The reason we need to do this mutably is because the query `(_ _ @last . ")" .)`, despite the liberal use
; of anchors, is broken (due to a bug in `tree-sitter`). Specifically, it will match both `b` and the
; following `,` in the tuple expression `(a, b,)`. This means we cannot set the attribute in this stanza
; (since overwriting attributes is not allowed) and so we instead write it to a mutable variable and set it
; later. Because the order in which the captures are returned results in `b` being matched before `,` this
; gives the correct behaviour.
[
(generator_expression
body: (_ _ @last . ")" .)
)
(list_comprehension
body: (_ _ @last . ")" .)
)
(set_comprehension
body: (_ _ @last . ")" .)
)
(dictionary_comprehension
body: (_ _ @last . ")" .)
)
] @genexpr
{
set @genexpr.body_end = (location-end @last)
}
[
(generator_expression)
(list_comprehension)
(set_comprehension)
(dictionary_comprehension)
] @genexpr
{
attr (@genexpr.expr) _location_end = @genexpr.body_end
attr (@genexpr.yield) _location_end = @genexpr.body_end
}
;;;;;; End of GeneratorExp (`(a for b in c if d)`)
;;;;;; Class statements
; A class definition
;
; class Foo(*bases, **keywords): body
;
; is turned into an actual assignment statement, with the class name as the left-hand side.
;
; Foo = $classexpr(name='Foo', bases, keywords, inner_scope=$class(name='Foo', body))
;
; (with a suitably magical definition of the `$` prefix).
;
; So we have to synthesize both the outer assignment, and also the two representatives of the class.
(class_definition
name: (identifier) @name
":" @colon
) @class
{
; To make it clearer that the outer node is an assignment, we create an alias for it.
let @class.assign = @class.node
; We reuse the identifier as the left hand side of the assignment.
let @class.assign_lhs = @name.node
; Synthesized nodes: the class_expr node, and the class node.
let @class.class_expr = (ast-node @class "ClassExpr")
let @class.inner_scope = (ast-node @class "Class")
; Setting up the outer assignment
edge @class.assign -> @class.assign_lhs
attr (@class.assign -> @class.assign_lhs) targets = 0
attr (@class.assign) value = @class.class_expr
attr (@class.assign) _location_end = (location-end @colon)
attr (@class.assign_lhs) ctx = "store"
let class_name = (source-text @name)
; The right hand side of the assignment, a `ClassExpr`.
attr (@class.class_expr) name = class_name
attr (@class.class_expr) inner_scope = @class.inner_scope
; `bases` will be set elsewhere
; `keywords` will be set elsewhere
attr (@class.class_expr) _location_end = (location-end @colon)
; The inner scope of the class_expr, a `Class`.
attr (@class.inner_scope) name = class_name
; body will be set in a separate stanza.
attr (@class.inner_scope) _location_end = (location-end @colon)
}
; Class.body
(class_definition
body: (block (_) @stmt)
) @class
{
edge @class.inner_scope -> @stmt.node
attr (@class.inner_scope -> @stmt.node) body = (named-child-index @stmt)
}
; Class.bases - using `(_ !value !name)` as a proxy for all non-keyword arguments.
; In particular, `keyword_argument` nodes have a `name` field, and `dictionary_splat`
; nodes have a `value` field.
(class_definition
superclasses: (argument_list element: (_ !value !name) @arg)
) @class
{
edge @class.class_expr -> @arg.node
attr (@class.class_expr -> @arg.node) bases = (named-child-index @arg)
}
; Class.keywords of the form `foo=bar`
(class_definition
superclasses: (argument_list element: (keyword_argument) @arg)
) @class
{
edge @class.class_expr -> @arg.node
attr (@class.class_expr -> @arg.node) keywords = (named-child-index @arg)
}
; Class.keywords of the form `**kwargs`
(class_definition
superclasses: (argument_list element: (dictionary_splat) @arg)
) @class
{
edge @class.class_expr -> @arg.node
attr (@class.class_expr -> @arg.node) keywords = (named-child-index @arg)
}
;;;;;; End of Class
;;;;;; Assign statements
; Assignment statements require a bit of interesting handling, since we represent a chained
; assignment such as `a = b = 5` as a single `Assign` node with multiple targets and a single
; right-hand side. This makes it somewhat complicated (but still doable) to determine the index of
; any single target in the resulting list.
;
; The way we handle this is by explicitly propagating two variables inwards. The first variable
; keeps track of the outermost node in a chain of assignments, and the second variable keeps track of
; the index of the left-hand side of the current assignment.
; Base case, for the outermost assignment we set the outermost node to this node, and the index to zero.
(expression_statement (assignment !type) @assign) @expr
{
let @assign.outermost_assignment = @assign.node
let @assign.target_index = 0
}
; Propagating the two variables inwards, increasing the index by one. Note that this depends on
; having the query match from the outside in -- if this evaluation order ever changes, this will break.
(assignment !type right: (assignment) @inner) @outer
{
let @inner.outermost_assignment = @outer.outermost_assignment
let @inner.target_index = (plus @outer.target_index 1)
}
; Finally, with the above variables set, we can -- for each assignment -- create an edge from the
; outermost assignment to it, and set its index to the index that we've calculated for this node.
(assignment !type left: (_) @target) @assign
{
edge @assign.outermost_assignment -> @target.node
attr (@assign.outermost_assignment -> @target.node) targets = @assign.target_index
attr (@target.node) ctx = "store"
}
; In addition to the above, we must ensure that the `value` attribute of the outermost assignment
; points to the _innermost_ right-hand side. We do this by first setting the `value` attribute for
; _all_ assignments...
(assignment !type right: (_) @value) @assign
{
attr (@assign.node) value = @value.node
attr (@value.node) ctx = "load"
}
; ... and then for assignments that are _inside_ other assigments, we use the `_skip_to` attribute
; to jump across the outer assignment.
;
; Thus, the outermost assignment's `value` will point to its right-hand side, but this one will (if
; it's an assignment itself) skip to _its_ right-hand side, and so on until we reach a right-hand side
; that is not an assignment.
(assignment !type right: (assignment right: (_) @inner) @outer)
{
attr (@outer.node) _skip_to = @inner.node
}
;;;;;; End of Assign
;;;;;; AnnAssign
(assignment
left: (_) @target
type: (type (expression) @type)
) @assign
{
attr (@assign.node) target = @target.node
attr (@target.node) ctx = "store"
attr (@assign.node) annotation = @type.node
attr (@type.node) ctx = "load"
}
(assignment
left: (_) @target
type: (_)
right: (_) @value
) @assign
{
attr (@assign.node) value = @value.node
attr (@value.node) ctx = "load"
}
;;;;;; End of AnnAssign
;;;;;; AugAssign
(augmented_assignment
left: (_) @left
operator: _ @op
right: (_) @right
) @augassign
{
let binop = (ast-node @augassign "BinOp")
attr (@augassign.node) operation = binop
attr (binop) left = @left.node
attr (@left.node) ctx = "load" ; yes, it really is "load".
attr (binop) op = (source-text @op)
attr (binop) right = @right.node
attr (@right.node) ctx = "load"
}
;;;;;; End of AugAssign
;;;;;; Global
(global_statement (identifier) @name) @global
{
edge @global.node -> @name.node
attr (@global.node -> @name.node) names = (named-child-index @name)
attr (@name.node) _is_literal = (source-text @name)
}
;;;;;; End of Global
;;;;;; Nonlocal
(nonlocal_statement (identifier) @name) @nonlocal
{
edge @nonlocal.node -> @name.node
attr (@nonlocal.node -> @name.node) names = (named-child-index @name)
attr (@name.node) _is_literal = (source-text @name)
}
;;;;;; End of Nonlocal
;;;;;; Import (`import ...`)
; `import j1.j2 as j3, j4, ...` becomes
;
; Import:
; names: [
; alias:
; value:
; ImportExpr:
; level: 0 # always 0 for absolute imports
; name: 'j1.j2'
; top: False
; asname:
; Name:
; variable: Variable('j3', None)
; ctx: Store
; alias:
; value:
; ImportExpr:
; level: 0 # always 0 for absolute imports
; name: 'j4'
; top: True
; asname:
; Name:
; variable: Variable('j4', None)
; ctx: Store
; ...
; ]
;
; from
;
; module
; import_statement
; name: aliased_import
; name: dotted_name
; identifier # j1
; identifier # j2
; alias: identifier j3
; name: dotted_name
; identifier # j4
;
; This means we have to hang our `alias` nodes off of the `dotted_name` and
; `aliased_import` nodes.
; Import.names
(import_statement name: (_) @name) @import
{
edge @import.node -> @name.node
attr (@import.node -> @name.node) names = (named-child-index @name)
}
; Imports without an explicit alias -- extract the root module name
(import_statement name: (dotted_name . (identifier) @first) @alias)
{
let import_expr = (ast-node @alias "ImportExpr")
attr (import_expr) level = 0
attr (import_expr) name = (source-text @alias)
attr (import_expr) top = #true
attr (@alias.node) value = import_expr
attr (@alias.node) asname = @first.node
attr (@first.node) ctx = "store"
}
; Not strictly needed (but the AST reconstruction will complain otherwise) we
; assign a context to each identifier in a dotted name (except the first part,
; which already gets one elsewhere).
(dotted_name (identifier) (identifier) @name)
{
attr (@name.node) ctx = "load"
}
; For dotted imports `a.b.c` the location for the `Name` corresponding to the
; `a` part covers the entire expression, so we explicitly match the final
; element and set the location appropriately. If there is only one element,
; this stanza doesn't fire, but in that case the location is actually correct
; already.
(import_statement
name: (dotted_name
.
(identifier) @first
(identifier) @last
.
)
)
{
attr (@first.node) _location_end = (location-end @last)
}
; Imports with an explicit alias
(import_statement
(aliased_import
name: (dotted_name . (identifier) @first) @name
alias: (identifier) @asname
) @alias
)
{
let import_expr = (ast-node @name "ImportExpr")
attr (import_expr) level = 0
attr (import_expr) name = (source-text @name)
attr (import_expr) top = #false
attr (@alias.node) value = import_expr
attr (@alias.node) asname = @asname.node
attr (@asname.node) ctx = "store"
attr (@first.node) ctx = "load"
}
;;;;;; End of Import (`import ...`)
;;;;;; Import (`from ... import ...`)
; Oh what a twisty mess these are. First, the prototypical layout of a
; `from some_module import x1 as y1, x2, ...` statement is as follows:
;
; Import:
; names: [
; alias:
; value:
; ImportMember:
; module:
; ImportExpr;
; level: