Merge pull request #17873 from github/tausbn/python-fix-generator-expression-locations

Python: Even more parser fixes

Author: tausbn

python/extractor/tests/parser/async-await.py

@@ -0,0 +1,8 @@
+async def foo():
+    await bar() + await baz()
+
+    async with foo() as bar, baz() as quux:
+        pass
+
+    async for spam in eggs:
+        pass

python/extractor/tests/parser/collections.py

@@ -35,3 +35,5 @@
   x, y,
 #comment
 )
+
+((z,))

python/extractor/tests/parser/comprehensions.py

@@ -65,3 +65,10 @@
     d for e in f if g # comment
     # comment
 ] # comment
+
+# Generator expression with comments
+(# comment
+    alpha # comment
+    for beta in gamma # comment
+    # comment
+)

python/extractor/tests/parser/strings.py

@@ -77,3 +77,12 @@
     b'\xc5\xe5'
 if 35:
     f"{x=}"
+if 36:
+    r"a\"a"
+if 37:
+    r'a\'a'
+if 38:
+    r'a\\'
+if 39:
+    r'a\
+    '

python/extractor/tsg-python/python.tsg

@@ -404,7 +404,7 @@
 
 ;;; GeneratorExp
 
-(generator_expression . "(" . (comment)* . (_) @start (_) @end . (comment)* . ")" .) @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)
@@ -416,13 +416,13 @@
     attr (@if.node) _location_end = (location-end @expr)
 }
 
-(generator_expression . "(" . (comment)* . (_) @start (for_in_clause) @child (_) @end . (comment)* . ")" .) @genexpr
+(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)* . (_) @start (for_in_clause) @end . (comment)* . ")" .) @genexpr
+(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)
@@ -863,7 +863,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)* . (_) @start  (_) @end . (comment)* . ")" .) @genexpr
+(generator_expression . "(" . (comment)* . (expression) @start  [(for_in_clause) (if_clause)] @end . (comment)* . ")" .) @genexpr
 {
     ; Synthesize the `genexpr` function
     let @genexpr.fun = (ast-node @genexpr "Function")
@@ -2661,6 +2661,14 @@
     let @with.first = @first.node
 }
 
+; Async status
+; NOTE: We only set the `is_async` field on the _first_ clause of the `with` statement,
+; as this is the behaviour of the old parser.
+(with_statement "async" "with" @with_keyword (with_clause . (with_item) @with))
+{
+    attr (@with.node) is_async = #true
+}
+
 (with_item
     value: (_) @value
 ) @with
@@ -3264,6 +3272,16 @@
     }
 }
 
+; Async status
+(function_definition "async" "def" @def_keyword) @funcdef
+{
+    let start = (location-start @def_keyword)
+    attr (@funcdef.function) is_async = #true
+    attr (@funcdef.node) _location_start = start
+    attr (@funcdef.function) _location_start = start
+    attr (@funcdef.funcexpr) _location_start = start
+}
+
 ;;; Decorators
 
 (decorated_definition
@@ -3478,5 +3496,9 @@
 
 [(tuple element: (_)) (tuple_pattern)] @tup
 {
-    attr (@tup.node) parenthesised = #true
+    ; In order to avoid writing to the `parenthesised` attribute twice, we only set it here
+    ; if the surrounding expression is not a `parenthesized_expression`.
+    if (not (instance-of (get-parent @tup) "parenthesized_expression")) {
+        attr (@tup.node) parenthesised = #true
+    }
 }

python/extractor/tsg-python/tsp/grammar.js

@@ -751,14 +751,14 @@ module.exports = grammar({
       $.comparison_operator,
       $.not_operator,
       $.boolean_operator,
-      $.await,
       $.lambda,
       $.primary_expression,
       $.conditional_expression,
       $.named_expression
     ),
 
     primary_expression: $ => choice(
+      $.await,
       $.binary_operator,
       $.identifier,
       $.keyword_identifier,
@@ -1202,7 +1202,7 @@ module.exports = grammar({
 
     await: $ => prec(PREC.unary, seq(
       'await',
-      $.expression
+      $.primary_expression
     )),
 
     comment: $ => token(seq('#', /.*/)),

python/extractor/tsg-python/tsp/src/grammar.json

@@ -1,5 +1,4 @@
 {
-  "$schema": "https://tree-sitter.github.io/tree-sitter/assets/schemas/grammar.schema.json",
   "name": "python",
   "word": "identifier",
   "rules": {
@@ -3843,10 +3842,6 @@
           "type": "SYMBOL",
           "name": "boolean_operator"
         },
-        {
-          "type": "SYMBOL",
-          "name": "await"
-        },
         {
           "type": "SYMBOL",
           "name": "lambda"
@@ -3868,6 +3863,10 @@
     "primary_expression": {
       "type": "CHOICE",
       "members": [
+        {
+          "type": "SYMBOL",
+          "name": "await"
+        },
         {
           "type": "SYMBOL",
           "name": "binary_operator"
@@ -6586,7 +6585,7 @@
           },
           {
             "type": "SYMBOL",
-            "name": "expression"
+            "name": "primary_expression"
           }
         ]
       }
@@ -6696,3 +6695,4 @@
     "parameter"
   ]
 }
+

python/extractor/tsg-python/tsp/src/node-types.json

@@ -115,10 +115,6 @@
     "type": "expression",
     "named": true,
     "subtypes": [
-      {
-        "type": "await",
-        "named": true
-      },
       {
         "type": "boolean_operator",
         "named": true
@@ -229,6 +225,10 @@
         "type": "attribute",
         "named": true
       },
+      {
+        "type": "await",
+        "named": true
+      },
       {
         "type": "binary_operator",
         "named": true
@@ -587,7 +587,7 @@
       "required": true,
       "types": [
         {
-          "type": "expression",
+          "type": "primary_expression",
           "named": true
         }
       ]
@@ -2691,7 +2691,6 @@
   {
     "type": "module",
     "named": true,
-    "root": true,
     "fields": {},
     "children": {
       "multiple": true,
@@ -3816,10 +3815,6 @@
     "type": ":=",
     "named": false
   },
-  {
-    "type": ";",
-    "named": false
-  },
   {
     "type": "<",
     "named": false
@@ -3876,10 +3871,6 @@
     "type": "[",
     "named": false
   },
-  {
-    "type": "\\",
-    "named": false
-  },
   {
     "type": "]",
     "named": false