support for star expression in lists/tuples

PyCQA · Oct 6, 2024 · 3c9d83c · 3c9d83c
1 parent eaa1062
commit 3c9d83c
Show file tree

Hide file tree

Showing 2 changed files with 190 additions and 2 deletions.
diff --git a/baron/grammator_data_structures.py b/baron/grammator_data_structures.py
@@ -101,13 +101,33 @@ def testlist_part_next(pack):
     def testlist_comp_empty(empty):
         return []
 
+    @pg.production("testlist_comp : star_expr comma star_expr")
+    def testlist_comp_two_star_star(pack):
+        (test, comma, test2) = pack
+        return [test, comma, test2]
+
+    @pg.production("testlist_comp : star_expr comma test")
+    def testlist_comp_two_star_test(pack):
+        (test, comma, test2) = pack
+        return [test, comma, test2]
+
+    @pg.production("testlist_comp : test comma star_expr")
+    def testlist_comp_two_test_star(pack):
+        (test, comma, test2) = pack
+        return [test, comma, test2]
+
     @pg.production("testlist_comp : test comma test")
-    def testlist_comp_two(pack):
+    def testlist_comp_two_test_test(pack):
         (test, comma, test2) = pack
         return [test, comma, test2]
 
+    @pg.production("testlist_comp : star_expr comma testlist_comp")
+    def testlist_comp_star_more(pack):
+        (test, comma, testlist_comp) = pack
+        return [test, comma] + testlist_comp
+
     @pg.production("testlist_comp : test comma testlist_comp")
-    def testlist_comp_more(pack):
+    def testlist_comp_test_more(pack):
         (test, comma, testlist_comp) = pack
         return [test, comma] + testlist_comp
 
@@ -132,11 +152,21 @@ def listmaker_one(pack):
         (test,) = pack
         return [test]
 
+    @pg.production("listmaker : star_expr")
+    def listmaker_one_star(pack):
+        (test,) = pack
+        return [test]
+
     @pg.production("listmaker : test comma listmaker")
     def listmaker_more(pack):
         (test, comma, listmaker) = pack
         return [test, comma] + listmaker
 
+    @pg.production("listmaker : star_expr comma listmaker")
+    def listmaker_more_start(pack):
+        (test, comma, listmaker) = pack
+        return [test, comma] + listmaker
+
     @pg.production("atom : LEFT_BRACKET dictmaker RIGHT_BRACKET")
     def dict(pack):
         (left_bracket, dictmaker, right_bracket,) = pack

diff --git a/tests/test_grammator_data_structures.py b/tests/test_grammator_data_structures.py
@@ -89,6 +89,95 @@ def test_tuple_one():
         }
     ])
 
+def test_tuple_one_star():
+    "( *a, )"
+    parse_simple([
+        ('LEFT_PARENTHESIS', '(', [], []),
+        ('STAR', '*', [], []),
+        ('NAME', 'a'),
+        ('COMMA', ',', [], [('SPACE', ' ')]),
+        ('RIGHT_PARENTHESIS', ')', [], []),
+    ], [
+        {
+            "with_parenthesis": True,
+            "first_formatting": [],
+            "second_formatting": [],
+            "third_formatting": [],
+            "fourth_formatting": [],
+            "type": "tuple",
+            "value": [
+                {
+                    "type": "star_expression",
+                    "formatting": [],
+                    "value": {
+                        "type": "name",
+                        "value": "a",
+                    }
+                },
+                {
+                    "type": "comma",
+                    "first_formatting": [],
+                    "second_formatting": [{"type": "space", "value": " "}],
+                }
+            ],
+        }
+    ])
+
+def test_tuple_more_star():
+    "( *a, b, *c )"
+    parse_simple([
+        ('LEFT_PARENTHESIS', '(', [], []),
+        ('STAR', '*', [], []),
+        ('NAME', 'a'),
+        ('COMMA', ',', [], [('SPACE', ' ')]),
+        ('NAME', 'b'),
+        ('COMMA', ',', [], [('SPACE', ' ')]),
+        ('STAR', '*', [], []),
+        ('NAME', 'c'),
+        ('RIGHT_PARENTHESIS', ')', [], []),
+    ], [
+        {
+            "with_parenthesis": True,
+            "first_formatting": [],
+            "second_formatting": [],
+            "third_formatting": [],
+            "fourth_formatting": [],
+            "type": "tuple",
+            "value": [
+                {
+                    "type": "star_expression",
+                    "formatting": [],
+                    "value": {
+                        "type": "name",
+                        "value": "a",
+                    }
+                },
+                {
+                    "type": "comma",
+                    "first_formatting": [],
+                    "second_formatting": [{"type": "space", "value": " "}],
+                },
+                {
+                    "type": "name",
+                    "value": "b",
+                },
+                {
+                    "type": "comma",
+                    "first_formatting": [],
+                    "second_formatting": [{"type": "space", "value": " "}],
+                },
+                {
+                    "type": "star_expression",
+                    "formatting": [],
+                    "value": {
+                        "type": "name",
+                        "value": "c",
+                    }
+                }
+            ],
+        }
+    ])
+
 
 def test_tuple_many():
     "(a, b, c)"
@@ -220,6 +309,74 @@ def test_list_more():
         }
     ])
 
+def test_list_star_one():
+    "[*a]"
+    parse_simple([
+        ('LEFT_SQUARE_BRACKET', '[', [], []),
+        ('STAR', '*', [], []),
+        ('NAME', 'a'),
+        ('RIGHT_SQUARE_BRACKET', ']', [], []),
+    ], [
+        {
+            "type": "list",
+            "first_formatting": [],
+            "second_formatting": [],
+            "third_formatting": [],
+            "fourth_formatting": [],
+            "value": [
+                {
+                    "formatting": [],
+                    "type": "star_expression",
+                    "value": {
+                        "type": "name",
+                        "value": "a",
+                    }
+                }
+            ],
+        }
+    ])
+
+def test_list_star_more():
+    "[*a, 'foo']"
+    parse_simple([
+        ('LEFT_SQUARE_BRACKET', '[', [], []),
+        ('STAR', '*', [], []),
+        ('NAME', 'a'),
+        ('COMMA', ',', [], [('SPACE', ' ')]),
+        ('STRING', "'foo'"),
+        ('RIGHT_SQUARE_BRACKET', ']', [], []),
+    ], [
+        {
+            "type": "list",
+            "first_formatting": [],
+            "second_formatting": [],
+            "third_formatting": [],
+            "fourth_formatting": [],
+            "value": [
+                {
+                    "type": "star_expression",
+                    "formatting": [],
+                    "value": {
+                        "type": "name",
+                        "value": "a",
+                    }
+                },
+                {
+                    "first_formatting": [],
+                    "second_formatting": [{"type": "space", "value": " "}],
+                    "type": "comma",
+                },
+                {
+                    "first_formatting": [],
+                    "second_formatting": [],
+                    "type": "string",
+                    "value": "'foo'",
+                }
+            ],
+        }
+    ])
+
+
 
 def test_dict_empty():
     "{ }"
@@ -301,6 +458,7 @@ def test_dict_one_double_star():
 )
 
 
+
 def test_dict_more_colon():
     "{a: b, b: c, c: d}"
     parse_simple([