SONARPY-1573: Fix parse error of unparenthesized walrus in set literals and comprehension (#2026)

* SONARPY-1573 parse set comprehension with unparenthesized walrus

* SONARPY-1573 adress PR feedback

Author: Seppli11

python-frontend/src/main/java/org/sonar/python/api/PythonGrammarBuilder.java

@@ -145,7 +145,7 @@ protected void grammar(LexerfulGrammarBuilder b) {
       b.sequence(
         b.firstOf(b.sequence(TEST, ":", TEST), b.sequence("**", EXPR)),
         b.firstOf(COMP_FOR, b.sequence(b.zeroOrMore(",", b.firstOf(b.sequence(TEST, ":", TEST), b.sequence("**", EXPR))), b.optional(",")))),
-      b.sequence(b.firstOf(TEST, STAR_EXPR), b.firstOf(COMP_FOR, b.sequence(b.zeroOrMore(",", b.firstOf(TEST, STAR_EXPR)), b.optional(","))))));
+      b.sequence(STAR_NAMED_EXPRESSION, b.firstOf(COMP_FOR, b.sequence(b.zeroOrMore(",", STAR_NAMED_EXPRESSION), b.optional(","))))));
 
     b.rule(ARGLIST).is(ARGUMENT, b.zeroOrMore(",", ARGUMENT), b.optional(","));
     b.rule(ARGUMENT).is(b.firstOf(

python-frontend/src/main/java/org/sonar/python/tree/PythonTreeMaker.java

@@ -1339,7 +1339,7 @@ private Expression dictOrSetLiteral(AstNode astNode) {
         Expression valueExpression = expression(dictOrSetMaker.getLastChild(PythonGrammar.TEST));
         return new DictCompExpressionImpl(lCurlyBrace, keyExpression, toPyToken(colon.getToken()), valueExpression, compFor, rCurlyBrace);
       } else {
-        Expression resultExpression = expression(dictOrSetMaker.getFirstChild(PythonGrammar.TEST, PythonGrammar.STAR_EXPR));
+        Expression resultExpression = expression(dictOrSetMaker.getFirstChild(PythonGrammar.STAR_NAMED_EXPRESSION));
         return new ComprehensionExpressionImpl(Tree.Kind.SET_COMPREHENSION, lCurlyBrace, resultExpression, compFor, rCurlyBrace);
       }
     }
@@ -1360,7 +1360,7 @@ private Expression dictOrSetLiteral(AstNode astNode) {
       }
       return new DictionaryLiteralImpl(lCurlyBrace, commas, dictionaryLiteralElements, rCurlyBrace);
     }
-    List<Expression> expressions = dictOrSetMaker.getChildren(PythonGrammar.TEST, PythonGrammar.STAR_EXPR).stream().map(this::expression).toList();
+    List<Expression> expressions = dictOrSetMaker.getChildren(PythonGrammar.STAR_NAMED_EXPRESSION).stream().map(this::expression).toList();
     return new SetLiteralImpl(lCurlyBrace, expressions, commas, rCurlyBrace);
   }
 

python-frontend/src/test/java/org/sonar/python/tree/PythonTreeMakerTest.java

@@ -2574,6 +2574,11 @@ void set_literal() {
     assertThat(tree.children()).hasSize(6).containsExactly(tree.lCurlyBrace(), tree.elements().get(0), tree.commas().get(0),
       tree.elements().get(1), tree.commas().get(1), tree.rCurlyBrace());
 
+    tree = (SetLiteral) parse("{ first := x, second := y }", treeMaker::expression);
+    assertThat(tree.elements()).hasSize(2);
+    assertThat(tree.children()).hasSize(5).containsExactly(tree.lCurlyBrace(), tree.elements().get(0), tree.commas().get(0),
+      tree.elements().get(1), tree.rCurlyBrace());
+
     tree = (SetLiteral) parse("{ *x }", treeMaker::expression);
     assertThat(tree.elements()).hasSize(1);
     element = tree.elements().iterator().next();
@@ -2592,8 +2597,21 @@ void set_comprehension() {
     assertThat(comprehension.children()).hasSize(4);
     assertThat(((Token) comprehension.children().get(0)).value()).isEqualTo("{");
     assertThat(((Token) comprehension.children().get(3)).value()).isEqualTo("}");
+  }
+
+  @Test
+  void set_comprehension_with_walrus_operator() {
+    setRootRule(PythonGrammar.TEST);
+    ComprehensionExpression comprehension =
+      (ComprehensionExpression) parse("{last := x-1 for x in [42, 43]}", treeMaker::expression);
     assertThat(comprehension.firstToken().value()).isEqualTo("{");
     assertThat(comprehension.lastToken().value()).isEqualTo("}");
+    assertThat(comprehension.getKind()).isEqualTo(Tree.Kind.SET_COMPREHENSION);
+    assertThat(comprehension.resultExpression().getKind()).isEqualTo(Kind.ASSIGNMENT_EXPRESSION);
+    assertThat(((AssignmentExpression) comprehension.resultExpression()).operator().value()).isEqualTo(":=");
+    assertThat(comprehension.children()).hasSize(4);
+    assertThat(((Token) comprehension.children().get(0)).value()).isEqualTo("{");
+    assertThat(((Token) comprehension.children().get(3)).value()).isEqualTo("}");
   }
 
   @Test

python-frontend/src/test/resources/parser/python/own/set_comprehension.py

@@ -0,0 +1,4 @@
+{x for x in range(7)}
+{(x4 := x**5) for x in range(7)}
+{x4 := x**5 for x in range(7)}
+{x := 1, 2, 3}