PY-48940: Allow unparenthesized assignment expressions within set literals, set comprehensions.

(cherry picked from commit 1ed44bb7a0bde2c7000923c2ed2a671f913eae8f)

IJ-CR-11891

GitOrigin-RevId: bb0ae8626e085e6c9223626137dbdd07e5c084b3

python/python-psi-impl/src/com/jetbrains/python/parsing/ExpressionParsing.java

@@ -23,6 +23,7 @@ import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.tree.TokenSet;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyTokenTypes;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import static com.jetbrains.python.PyPsiBundle.message;
@@ -342,8 +343,13 @@ public class ExpressionParsing extends Parsing {
       return;
     }
 
+    parseDictOrSetTail(expr, false);
+  }
+
+  private void parseDictOrSetTail(@NotNull SyntaxTreeBuilder.Marker expr, boolean isDict) {
     final SyntaxTreeBuilder.Marker firstExprMarker = myBuilder.mark();
-    if (!parseSingleExpression(false)) {
+
+    if (isDict && !parseSingleExpression(false) || !isDict && !parseNamedTestExpression(false, false)) {
       myBuilder.error(message("PARSE.expected.expression"));
       firstExprMarker.drop();
       expr.done(PyElementTypes.DICT_LITERAL_EXPRESSION);
@@ -351,6 +357,11 @@ public class ExpressionParsing extends Parsing {
     }
 
     if (matchToken(PyTokenTypes.COLON)) {
+      if (!isDict) {
+        firstExprMarker.rollbackTo();
+        parseDictOrSetTail(expr, true);
+        return;
+      }
       parseDictLiteralTail(expr, firstExprMarker);
     }
     else if (atToken(PyTokenTypes.COMMA) || atToken(PyTokenTypes.RBRACE)) {
@@ -424,7 +435,7 @@ public class ExpressionParsing extends Parsing {
   private void parseSetLiteralTail(SyntaxTreeBuilder.Marker startMarker) {
     while (myBuilder.getTokenType() != PyTokenTypes.RBRACE) {
       checkMatches(PyTokenTypes.COMMA, message("PARSE.expected.comma"));
-      if (!parseSingleExpression(false)) {
+      if (!parseNamedTestExpression(false, false)) {
         break;
       }
     }

python/testData/psi/AssignmentExpressionsInSet.py

@@ -0,0 +1,7 @@
+old_set = {(a := 1), (b := 1)}    # valid
+new_set = {a := 1, b := 2, c := 3}  # valid
+
+my_list = [1, 2, 3]
+set_comp_old = {(a := my_list[i]) for i in my_list if (k := i) > 0}    # valid
+set_comp_new = {b := my_list[j] for j in my_list if (k := j) > 0}     # valid
+set_comp_new_invalid = {b := my_list[j] for j in my_list if k := True}     # invalid

python/testData/psi/AssignmentExpressionsInSet.txt

@@ -0,0 +1,261 @@
+PyFile:AssignmentExpressionsInSet.py
+  PyAssignmentStatement
+    PyTargetExpression: old_set
+      PsiElement(Py:IDENTIFIER)('old_set')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:EQ)('=')
+    PsiWhiteSpace(' ')
+    PySetLiteralExpression
+      PsiElement(Py:LBRACE)('{')
+      PyParenthesizedExpression
+        PsiElement(Py:LPAR)('(')
+        PyAssignmentExpression
+          PyTargetExpression: a
+            PsiElement(Py:IDENTIFIER)('a')
+          PsiWhiteSpace(' ')
+          PsiElement(Py:COLONEQ)(':=')
+          PsiWhiteSpace(' ')
+          PyNumericLiteralExpression
+            PsiElement(Py:INTEGER_LITERAL)('1')
+        PsiElement(Py:RPAR)(')')
+      PsiElement(Py:COMMA)(',')
+      PsiWhiteSpace(' ')
+      PyParenthesizedExpression
+        PsiElement(Py:LPAR)('(')
+        PyAssignmentExpression
+          PyTargetExpression: b
+            PsiElement(Py:IDENTIFIER)('b')
+          PsiWhiteSpace(' ')
+          PsiElement(Py:COLONEQ)(':=')
+          PsiWhiteSpace(' ')
+          PyNumericLiteralExpression
+            PsiElement(Py:INTEGER_LITERAL)('1')
+        PsiElement(Py:RPAR)(')')
+      PsiElement(Py:RBRACE)('}')
+    PsiWhiteSpace('    ')
+    PsiComment(Py:END_OF_LINE_COMMENT)('# valid')
+  PsiWhiteSpace('\n')
+  PyAssignmentStatement
+    PyTargetExpression: new_set
+      PsiElement(Py:IDENTIFIER)('new_set')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:EQ)('=')
+    PsiWhiteSpace(' ')
+    PySetLiteralExpression
+      PsiElement(Py:LBRACE)('{')
+      PyAssignmentExpression
+        PyTargetExpression: a
+          PsiElement(Py:IDENTIFIER)('a')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:COLONEQ)(':=')
+        PsiWhiteSpace(' ')
+        PyNumericLiteralExpression
+          PsiElement(Py:INTEGER_LITERAL)('1')
+      PsiElement(Py:COMMA)(',')
+      PsiWhiteSpace(' ')
+      PyAssignmentExpression
+        PyTargetExpression: b
+          PsiElement(Py:IDENTIFIER)('b')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:COLONEQ)(':=')
+        PsiWhiteSpace(' ')
+        PyNumericLiteralExpression
+          PsiElement(Py:INTEGER_LITERAL)('2')
+      PsiElement(Py:COMMA)(',')
+      PsiWhiteSpace(' ')
+      PyAssignmentExpression
+        PyTargetExpression: c
+          PsiElement(Py:IDENTIFIER)('c')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:COLONEQ)(':=')
+        PsiWhiteSpace(' ')
+        PyNumericLiteralExpression
+          PsiElement(Py:INTEGER_LITERAL)('3')
+      PsiElement(Py:RBRACE)('}')
+    PsiWhiteSpace('  ')
+    PsiComment(Py:END_OF_LINE_COMMENT)('# valid')
+  PsiWhiteSpace('\n\n')
+  PyAssignmentStatement
+    PyTargetExpression: my_list
+      PsiElement(Py:IDENTIFIER)('my_list')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:EQ)('=')
+    PsiWhiteSpace(' ')
+    PyListLiteralExpression
+      PsiElement(Py:LBRACKET)('[')
+      PyNumericLiteralExpression
+        PsiElement(Py:INTEGER_LITERAL)('1')
+      PsiElement(Py:COMMA)(',')
+      PsiWhiteSpace(' ')
+      PyNumericLiteralExpression
+        PsiElement(Py:INTEGER_LITERAL)('2')
+      PsiElement(Py:COMMA)(',')
+      PsiWhiteSpace(' ')
+      PyNumericLiteralExpression
+        PsiElement(Py:INTEGER_LITERAL)('3')
+      PsiElement(Py:RBRACKET)(']')
+  PsiWhiteSpace('\n')
+  PyAssignmentStatement
+    PyTargetExpression: set_comp_old
+      PsiElement(Py:IDENTIFIER)('set_comp_old')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:EQ)('=')
+    PsiWhiteSpace(' ')
+    PySetCompExpression
+      PsiElement(Py:LBRACE)('{')
+      PyParenthesizedExpression
+        PsiElement(Py:LPAR)('(')
+        PyAssignmentExpression
+          PyTargetExpression: a
+            PsiElement(Py:IDENTIFIER)('a')
+          PsiWhiteSpace(' ')
+          PsiElement(Py:COLONEQ)(':=')
+          PsiWhiteSpace(' ')
+          PySubscriptionExpression
+            PyReferenceExpression: my_list
+              PsiElement(Py:IDENTIFIER)('my_list')
+            PsiElement(Py:LBRACKET)('[')
+            PyReferenceExpression: i
+              PsiElement(Py:IDENTIFIER)('i')
+            PsiElement(Py:RBRACKET)(']')
+        PsiElement(Py:RPAR)(')')
+      PsiWhiteSpace(' ')
+      PsiElement(Py:FOR_KEYWORD)('for')
+      PsiWhiteSpace(' ')
+      PyTargetExpression: i
+        PsiElement(Py:IDENTIFIER)('i')
+      PsiWhiteSpace(' ')
+      PsiElement(Py:IN_KEYWORD)('in')
+      PsiWhiteSpace(' ')
+      PyReferenceExpression: my_list
+        PsiElement(Py:IDENTIFIER)('my_list')
+      PsiWhiteSpace(' ')
+      PsiElement(Py:IF_KEYWORD)('if')
+      PsiWhiteSpace(' ')
+      PyBinaryExpression
+        PyParenthesizedExpression
+          PsiElement(Py:LPAR)('(')
+          PyAssignmentExpression
+            PyTargetExpression: k
+              PsiElement(Py:IDENTIFIER)('k')
+            PsiWhiteSpace(' ')
+            PsiElement(Py:COLONEQ)(':=')
+            PsiWhiteSpace(' ')
+            PyReferenceExpression: i
+              PsiElement(Py:IDENTIFIER)('i')
+          PsiElement(Py:RPAR)(')')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:GT)('>')
+        PsiWhiteSpace(' ')
+        PyNumericLiteralExpression
+          PsiElement(Py:INTEGER_LITERAL)('0')
+      PsiElement(Py:RBRACE)('}')
+    PsiWhiteSpace('    ')
+    PsiComment(Py:END_OF_LINE_COMMENT)('# valid')
+  PsiWhiteSpace('\n')
+  PyAssignmentStatement
+    PyTargetExpression: set_comp_new
+      PsiElement(Py:IDENTIFIER)('set_comp_new')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:EQ)('=')
+    PsiWhiteSpace(' ')
+    PySetCompExpression
+      PsiElement(Py:LBRACE)('{')
+      PyAssignmentExpression
+        PyTargetExpression: b
+          PsiElement(Py:IDENTIFIER)('b')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:COLONEQ)(':=')
+        PsiWhiteSpace(' ')
+        PySubscriptionExpression
+          PyReferenceExpression: my_list
+            PsiElement(Py:IDENTIFIER)('my_list')
+          PsiElement(Py:LBRACKET)('[')
+          PyReferenceExpression: j
+            PsiElement(Py:IDENTIFIER)('j')
+          PsiElement(Py:RBRACKET)(']')
+      PsiWhiteSpace(' ')
+      PsiElement(Py:FOR_KEYWORD)('for')
+      PsiWhiteSpace(' ')
+      PyTargetExpression: j
+        PsiElement(Py:IDENTIFIER)('j')
+      PsiWhiteSpace(' ')
+      PsiElement(Py:IN_KEYWORD)('in')
+      PsiWhiteSpace(' ')
+      PyReferenceExpression: my_list
+        PsiElement(Py:IDENTIFIER)('my_list')
+      PsiWhiteSpace(' ')
+      PsiElement(Py:IF_KEYWORD)('if')
+      PsiWhiteSpace(' ')
+      PyBinaryExpression
+        PyParenthesizedExpression
+          PsiElement(Py:LPAR)('(')
+          PyAssignmentExpression
+            PyTargetExpression: k
+              PsiElement(Py:IDENTIFIER)('k')
+            PsiWhiteSpace(' ')
+            PsiElement(Py:COLONEQ)(':=')
+            PsiWhiteSpace(' ')
+            PyReferenceExpression: j
+              PsiElement(Py:IDENTIFIER)('j')
+          PsiElement(Py:RPAR)(')')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:GT)('>')
+        PsiWhiteSpace(' ')
+        PyNumericLiteralExpression
+          PsiElement(Py:INTEGER_LITERAL)('0')
+      PsiElement(Py:RBRACE)('}')
+    PsiWhiteSpace('     ')
+    PsiComment(Py:END_OF_LINE_COMMENT)('# valid')
+  PsiWhiteSpace('\n')
+  PyAssignmentStatement
+    PyTargetExpression: set_comp_new_invalid
+      PsiElement(Py:IDENTIFIER)('set_comp_new_invalid')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:EQ)('=')
+    PsiWhiteSpace(' ')
+    PyAssignmentExpression
+      PySetCompExpression
+        PsiElement(Py:LBRACE)('{')
+        PyAssignmentExpression
+          PyTargetExpression: b
+            PsiElement(Py:IDENTIFIER)('b')
+          PsiWhiteSpace(' ')
+          PsiElement(Py:COLONEQ)(':=')
+          PsiWhiteSpace(' ')
+          PySubscriptionExpression
+            PyReferenceExpression: my_list
+              PsiElement(Py:IDENTIFIER)('my_list')
+            PsiElement(Py:LBRACKET)('[')
+            PyReferenceExpression: j
+              PsiElement(Py:IDENTIFIER)('j')
+            PsiElement(Py:RBRACKET)(']')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:FOR_KEYWORD)('for')
+        PsiWhiteSpace(' ')
+        PyTargetExpression: j
+          PsiElement(Py:IDENTIFIER)('j')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:IN_KEYWORD)('in')
+        PsiWhiteSpace(' ')
+        PyReferenceExpression: my_list
+          PsiElement(Py:IDENTIFIER)('my_list')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:IF_KEYWORD)('if')
+        PsiWhiteSpace(' ')
+        PyReferenceExpression: k
+          PsiElement(Py:IDENTIFIER)('k')
+        PsiErrorElement:']' or 'for' expected
+          <empty list>
+      PsiWhiteSpace(' ')
+      PsiElement(Py:COLONEQ)(':=')
+      PsiWhiteSpace(' ')
+      PyBoolLiteralExpression
+        PsiElement(Py:TRUE_KEYWORD)('True')
+    PsiErrorElement:Expression expected
+      <empty list>
+  PsiElement(Py:RBRACE)('}')
+  PsiErrorElement:Statement expected, found Py:RBRACE
+    <empty list>
+  PsiWhiteSpace('     ')
+  PsiComment(Py:END_OF_LINE_COMMENT)('# invalid')

python/testData/psi/InvalidNonParenthesizedAssignmentExpressions.py

@@ -14,10 +14,8 @@ def foo(answer: (p := 42) = 5):  # Valid, but probably never useful
 (lambda: x := 1)  # INVALID
 lambda: (x := 1)  # Valid, but unlikely to be useful
 
-result_set = {a := 1}  # INVALID
-result_set = {(a := 1)}
-
 result_dict = {a := 1 : b := 2}  # INVALID
+result_dict = {a := 1 : (b := 2)}  # INVALID
 result_dict = {(a := 1) : (b := 2)}
 
 assert a := 1  # INVALID

python/testData/psi/InvalidNonParenthesizedAssignmentExpressions.txt

@@ -259,53 +259,6 @@ PyFile:InvalidNonParenthesizedAssignmentExpressions.py
     PsiWhiteSpace('  ')
     PsiComment(Py:END_OF_LINE_COMMENT)('# Valid, but unlikely to be useful')
   PsiWhiteSpace('\n\n')
-  PyAssignmentStatement
-    PyTargetExpression: result_set
-      PsiElement(Py:IDENTIFIER)('result_set')
-    PsiWhiteSpace(' ')
-    PsiElement(Py:EQ)('=')
-    PsiWhiteSpace(' ')
-    PyAssignmentExpression
-      PyDictLiteralExpression
-        PsiElement(Py:LBRACE)('{')
-        PyReferenceExpression: a
-          PsiElement(Py:IDENTIFIER)('a')
-        PsiErrorElement:Expression expected
-          <empty list>
-      PsiWhiteSpace(' ')
-      PsiElement(Py:COLONEQ)(':=')
-      PsiWhiteSpace(' ')
-      PyNumericLiteralExpression
-        PsiElement(Py:INTEGER_LITERAL)('1')
-    PsiErrorElement:Expression expected
-      <empty list>
-  PsiElement(Py:RBRACE)('}')
-  PsiErrorElement:Statement expected, found Py:RBRACE
-    <empty list>
-  PsiWhiteSpace('  ')
-  PsiComment(Py:END_OF_LINE_COMMENT)('# INVALID')
-  PsiWhiteSpace('\n')
-  PyAssignmentStatement
-    PyTargetExpression: result_set
-      PsiElement(Py:IDENTIFIER)('result_set')
-    PsiWhiteSpace(' ')
-    PsiElement(Py:EQ)('=')
-    PsiWhiteSpace(' ')
-    PySetLiteralExpression
-      PsiElement(Py:LBRACE)('{')
-      PyParenthesizedExpression
-        PsiElement(Py:LPAR)('(')
-        PyAssignmentExpression
-          PyTargetExpression: a
-            PsiElement(Py:IDENTIFIER)('a')
-          PsiWhiteSpace(' ')
-          PsiElement(Py:COLONEQ)(':=')
-          PsiWhiteSpace(' ')
-          PyNumericLiteralExpression
-            PsiElement(Py:INTEGER_LITERAL)('1')
-        PsiElement(Py:RPAR)(')')
-      PsiElement(Py:RBRACE)('}')
-  PsiWhiteSpace('\n\n')
   PyAssignmentStatement
     PyTargetExpression: result_dict
       PsiElement(Py:IDENTIFIER)('result_dict')
@@ -348,6 +301,51 @@ PyFile:InvalidNonParenthesizedAssignmentExpressions.py
   PsiWhiteSpace('  ')
   PsiComment(Py:END_OF_LINE_COMMENT)('# INVALID')
   PsiWhiteSpace('\n')
+  PyAssignmentStatement
+    PyTargetExpression: result_dict
+      PsiElement(Py:IDENTIFIER)('result_dict')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:EQ)('=')
+    PsiWhiteSpace(' ')
+    PyAssignmentExpression
+      PyDictLiteralExpression
+        PsiElement(Py:LBRACE)('{')
+        PyReferenceExpression: a
+          PsiElement(Py:IDENTIFIER)('a')
+        PsiErrorElement:Expression expected
+          <empty list>
+      PsiWhiteSpace(' ')
+      PsiElement(Py:COLONEQ)(':=')
+      PsiWhiteSpace(' ')
+      PyNumericLiteralExpression
+        PsiElement(Py:INTEGER_LITERAL)('1')
+    PsiErrorElement:Expression expected
+      <empty list>
+  PsiWhiteSpace(' ')
+  PsiElement(Py:COLON)(':')
+  PsiErrorElement:Statement expected, found Py:COLON
+    <empty list>
+  PsiWhiteSpace(' ')
+  PyExpressionStatement
+    PyParenthesizedExpression
+      PsiElement(Py:LPAR)('(')
+      PyAssignmentExpression
+        PyTargetExpression: b
+          PsiElement(Py:IDENTIFIER)('b')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:COLONEQ)(':=')
+        PsiWhiteSpace(' ')
+        PyNumericLiteralExpression
+          PsiElement(Py:INTEGER_LITERAL)('2')
+      PsiElement(Py:RPAR)(')')
+    PsiErrorElement:End of statement expected
+      <empty list>
+  PsiElement(Py:RBRACE)('}')
+  PsiErrorElement:Statement expected, found Py:RBRACE
+    <empty list>
+  PsiWhiteSpace('  ')
+  PsiComment(Py:END_OF_LINE_COMMENT)('# INVALID')
+  PsiWhiteSpace('\n')
   PyAssignmentStatement
     PyTargetExpression: result_dict
       PsiElement(Py:IDENTIFIER)('result_dict')
@@ -433,25 +431,22 @@ PyFile:InvalidNonParenthesizedAssignmentExpressions.py
         PsiElement(Py:INTEGER_LITERAL)('2')
       PsiElement(Py:RBRACKET)(']')
   PsiWhiteSpace('\n')
-  PyAssignmentExpression
+  PyExpressionStatement
     PySubscriptionExpression
       PyReferenceExpression: l
         PsiElement(Py:IDENTIFIER)('l')
       PsiElement(Py:LBRACKET)('[')
-      PyReferenceExpression: a
+      PyAssignmentExpression
-        PsiElement(Py:IDENTIFIER)('a')
+        PyTargetExpression: a
-      PsiErrorElement:']' expected
+          PsiElement(Py:IDENTIFIER)('a')
-        <empty list>
+        PsiWhiteSpace(' ')
-    PsiWhiteSpace(' ')
+        PsiElement(Py:COLONEQ)(':=')
-    PsiElement(Py:COLONEQ)(':=')
+        PsiWhiteSpace(' ')
-    PsiWhiteSpace(' ')
+        PyNumericLiteralExpression
-    PyNumericLiteralExpression
+          PsiElement(Py:INTEGER_LITERAL)('0')
-      PsiElement(Py:INTEGER_LITERAL)('0')
+      PsiElement(Py:RBRACKET)(']')
-  PsiElement(Py:RBRACKET)(']')
+    PsiWhiteSpace('  ')
-  PsiErrorElement:Statement expected, found Py:IDENTIFIER
+    PsiComment(Py:END_OF_LINE_COMMENT)('# INVALID')
-    <empty list>
-  PsiWhiteSpace('  ')
-  PsiComment(Py:END_OF_LINE_COMMENT)('# INVALID')
   PsiWhiteSpace('\n')
   PyExpressionStatement
     PySubscriptionExpression

python/testSrc/com/jetbrains/python/parsing/PythonParsingTest.java

@@ -1186,6 +1186,11 @@ public class PythonParsingTest extends ParsingTestCase {
     doTest(LanguageLevel.PYTHON310);
   }
 
+  // PY-48940
+  public void testAssignmentExpressionsInSet() {
+    doTest(LanguageLevel.getLatest());
+  }
+
   public void doTest() {
     doTest(LanguageLevel.PYTHON26);
   }