diff --git a/python/extractor/tsg-python/python.tsg b/python/extractor/tsg-python/python.tsg index 93d6e95a344..e0aadd5432c 100644 --- a/python/extractor/tsg-python/python.tsg +++ b/python/extractor/tsg-python/python.tsg @@ -403,7 +403,7 @@ ;;; GeneratorExp -(generator_expression . "(" . (comment)* . (expression) @start [(for_in_clause) (if_clause)] @end . (comment)* . ")" .) @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) @@ -415,13 +415,13 @@ 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 +(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) @start (for_in_clause) @end . (comment)* . ")" .) @_genexpr +(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) @@ -824,6 +824,29 @@ attr (@genexpr.arg_use) ctx = "load" } +; DictComp with unpacking (PEP 798): `{**d for d in dicts}` +(dictionary_comprehension + body: (dictionary_splat) +) @genexpr +{ + let @genexpr.fun = (ast-node @genexpr "Function") + attr (@genexpr.node) function = @genexpr.fun + attr (@genexpr.fun) name = "dictcomp" + + 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 + + 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)`) @@ -862,7 +885,7 @@ ; 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 +(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") @@ -1034,12 +1057,25 @@ ; 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 + (generator_expression body: (expression) @body) @genexpr + (list_comprehension body: (expression) @body) @genexpr + (set_comprehension body: (expression) @body) @genexpr ] { let @genexpr.result = @body.node + let @genexpr.yield_kind = "Yield" +} + +; For starred comprehensions (PEP 798), the result is the inner expression (not the Starred +; wrapper), and we use `yield from` instead of `yield` to represent the unpacking semantics. +[ + (generator_expression body: (list_splat (expression) @inner) @_body) @genexpr + (list_comprehension body: (list_splat (expression) @inner) @_body) @genexpr + (set_comprehension body: (list_splat (expression) @inner) @_body) @genexpr +] +{ + let @genexpr.result = @inner.node + let @genexpr.yield_kind = "YieldFrom" } ; For dict comprehensions, we build an explicit tuple using the key and value pair. @@ -1052,13 +1088,31 @@ { let tuple = (ast-node @body "Tuple") edge tuple -> @key.node - attr (tuple -> @key.node) elts = 1 + attr (tuple -> @key.node) elts = 0 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 -> @value.node) elts = 1 attr (tuple) ctx = "load" let @genexpr.result = tuple + let @genexpr.yield_kind = "Yield" +} + +; For dict comprehensions with unpacking (PEP 798), `{**d for d in dicts}` desugars to +; `yield from d.items()` to produce (key, value) tuples consistent with the regular dict comp model. +(dictionary_comprehension + body: (dictionary_splat (expression) @inner) @_body +) @genexpr +{ + ; Synthesize `d.items()`: Attribute(value=d, attr='items') then Call(func=attr) + let attr = (ast-node @inner "Attribute") + attr (attr) value = @inner.node + attr (attr) attr = "items" + attr (attr) ctx = "load" + + let call = (ast-node @inner "Call") + attr (call) func = attr + + let @genexpr.result = call + let @genexpr.yield_kind = "YieldFrom" } ; For the final clause, we need to hook it up with the rest of the expression. @@ -1094,7 +1148,7 @@ let last = (get-last-element @last_candidates) let expr = (ast-node @body "Expr") - let yield = (ast-node @body "Yield") + let yield = (ast-node @body @genexpr.yield_kind) let @genexpr.expr = expr let @genexpr.yield = yield