Don't redundantly nest StarredElement inside another Element (#624)

Instagram · Jan 23, 2022 · 2b2b25b · 2b2b25b
1 parent d35b6a5
commit 2b2b25b
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 3 deletions.
diff --git a/libcst/_nodes/tests/test_tuple.py b/libcst/_nodes/tests/test_tuple.py
@@ -91,6 +91,47 @@ class TupleTest(CSTNodeTest):
                 "parser": parse_expression,
                 "expected_position": CodeRange((1, 1), (1, 11)),
             },
+            # top-level two-element tuple, with one being starred
+            {
+                "node": cst.SimpleStatementLine(
+                    body=[
+                        cst.Expr(
+                            value=cst.Tuple(
+                                [
+                                    cst.Element(cst.Name("one"), comma=cst.Comma()),
+                                    cst.StarredElement(cst.Name("two")),
+                                ],
+                                lpar=[],
+                                rpar=[],
+                            )
+                        )
+                    ]
+                ),
+                "code": "one,*two\n",
+                "parser": parse_statement,
+            },
+            # top-level three-element tuple, start/end is starred
+            {
+                "node": cst.SimpleStatementLine(
+                    body=[
+                        cst.Expr(
+                            value=cst.Tuple(
+                                [
+                                    cst.StarredElement(
+                                        cst.Name("one"), comma=cst.Comma()
+                                    ),
+                                    cst.Element(cst.Name("two"), comma=cst.Comma()),
+                                    cst.StarredElement(cst.Name("three")),
+                                ],
+                                lpar=[],
+                                rpar=[],
+                            )
+                        )
+                    ]
+                ),
+                "code": "*one,two,*three\n",
+                "parser": parse_statement,
+            },
             # missing spaces around tuple, okay with parenthesis
             {
                 "node": cst.For(

diff --git a/native/libcst/src/parser/grammar.rs b/native/libcst/src/parser/grammar.rs
@@ -2174,9 +2174,12 @@ fn make_assignment<'a>(
 }
 
 fn expr_to_element(expr: Expression) -> Element {
-    Element::Simple {
-        value: expr,
-        comma: Default::default(),
+    match expr {
+        Expression::StarredElement(inner_expr) => Element::Starred(inner_expr),
+        _ => Element::Simple {
+            value: expr,
+            comma: Default::default(),
+        },
     }
 }