Improve parsing of for-from loops

b0o · Dec 8, 2024 · a1c9c26 · a1c9c26
1 parent 0fe3849
commit a1c9c26
Show file tree

Hide file tree

Showing 7 changed files with 229,300 additions and 233,638 deletions.
diff --git a/grammar.js b/grammar.js
@@ -780,19 +780,41 @@ module.exports = grammar(Python, {
         ),
       ),
 
-    for_statement: $ =>
+
+    for_in_loop: $ =>
       seq(
-        optional("async"),
-        "for",
-        optional(seq(
-          field("left", $._left_hand_side),
-          choice("in", "from"),
-        )),
+        field("left", $._left_hand_side),
+        "in",
         field("right", $._expressions),
+      ),
+
+    for_from_relation: $ =>
+      choice("<", "<=", ">", ">="),
+
+    for_from_loop: $ =>
+      seq(
+        optional(
+          seq(
+            field("left", $.identifier),
+            "from",
+          ),
+        ),
+        field("lower", choice($.primary_expression)),
+        $.for_from_relation,
+        field("var", $.identifier),
+        $.for_from_relation,
+        field("upper", choice($.primary_expression)),
         optional(seq(
           "by",
           $._expressions,
         )),
+      ),
+
+    for_statement: $ =>
+      seq(
+        optional("async"),
+        "for",
+        choice($.for_from_loop, $.for_in_loop),
         ":",
         field("body", $._suite),
         field("alternative", optional($.else_clause)),

diff --git a/queries/locals.scm b/queries/locals.scm
@@ -65,15 +65,15 @@
 
 ; Loops
 ; not a scope!
-(for_statement
+(for_in_loop
   left: (pattern_list
     (identifier) @local.definition))
 
-(for_statement
+(for_in_loop
   left: (tuple_pattern
     (identifier) @local.definition))
 
-(for_statement
+(for_in_loop
   left: (identifier) @local.definition)
 
 ; not a scope!

diff --git a/src/grammar.json b/src/grammar.json