PY-89217 PEP-798 - generator expression with unpacking as call argument

GitOrigin-RevId: 46422f3e5c7188fcb28f8624ef5f058bd2b10728

python/python-parser/src/com/jetbrains/python/parsing/ExpressionParsing.java

@@ -588,12 +588,19 @@ public class ExpressionParsing extends Parsing {
         }
       }
       if (myBuilder.getTokenType() == PyTokenTypes.MULT || myBuilder.getTokenType() == PyTokenTypes.EXP) {
+        final boolean singleStar = myBuilder.getTokenType() == PyTokenTypes.MULT;
         final SyntaxTreeBuilder.Marker starArgMarker = myBuilder.mark();
         myBuilder.advanceLexer();
         if (!parseSingleExpression(false)) {
           myBuilder.error(message("PARSE.expected.expression"));
         }
-        starArgMarker.done(PyElementTypes.STAR_ARGUMENT_EXPRESSION);
+        if (singleStar && argNumber == 1 && genexpr != null && atForOrAsyncFor()) {
+          // PEP 798: a generator expression with star unpacking passed as the sole argument, e.g. foo(*x for x in y)
+          starArgMarker.done(PyElementTypes.STAR_EXPRESSION);
+        }
+        else {
+          starArgMarker.done(PyElementTypes.STAR_ARGUMENT_EXPRESSION);
+        }
       }
       else {
         if (isIdentifier(myBuilder)) {

python/testData/inspections/PyCompatibilityInspection/unpackingInComprehensions.py

@@ -2,3 +2,4 @@
 {<warning descr="Python version 3.14 does not support unpacking in comprehensions">*it</warning> for it in its}
 {<warning descr="Python version 3.14 does not support unpacking in comprehensions">**d</warning> for d in dicts}
 (<warning descr="Python version 3.14 does not support unpacking in comprehensions">*it</warning> for it in its)
+foo(<warning descr="Python version 3.14 does not support unpacking in comprehensions">*it</warning> for it in its)

python/testData/psi/UnpackingInComprehensions.py

@@ -2,3 +2,4 @@
 {*it for it in its}
 {**d for d in dicts}
 (*it for it in its)
+foo(*it for it in its)

python/testData/psi/UnpackingInComprehensions.txt

@@ -74,3 +74,26 @@ PyFile:UnpackingInComprehensions.py
       PyReferenceExpression: its
         PsiElement(Py:IDENTIFIER)('its')
       PsiElement(Py:RPAR)(')')
+  PsiWhiteSpace('\n')
+  PyExpressionStatement
+    PyCallExpression: foo
+      PyReferenceExpression: foo
+        PsiElement(Py:IDENTIFIER)('foo')
+      PyArgumentList
+        PsiElement(Py:LPAR)('(')
+        PyGeneratorExpression
+          PyStarExpression
+            PsiElement(Py:MULT)('*')
+            PyReferenceExpression: it
+              PsiElement(Py:IDENTIFIER)('it')
+          PsiWhiteSpace(' ')
+          PsiElement(Py:FOR_KEYWORD)('for')
+          PsiWhiteSpace(' ')
+          PyTargetExpression: it
+            PsiElement(Py:IDENTIFIER)('it')
+          PsiWhiteSpace(' ')
+          PsiElement(Py:IN_KEYWORD)('in')
+          PsiWhiteSpace(' ')
+          PyReferenceExpression: its
+            PsiElement(Py:IDENTIFIER)('its')
+        PsiElement(Py:RPAR)(')')

python/testSrc/com/jetbrains/python/types/PyComprehensionAndIteratorTypeTest.kt

@@ -76,6 +76,14 @@ class PyComprehensionAndIteratorTypeTest : PyCodeInsightTestCase() {
       #   └ TYPE dict[str, int]
       """)
 
+    // PEP 798: a bare generator expression with star unpacking passed as the sole call argument
+    @Test
+    fun `unpacking in generator expression call argument type`() = test(TestOptions(assertRecursionPrevention = false), """
+      def f(its: list[list[int]]):
+          expr = list(*it for it in its)
+      #   └ TYPE list[int]
+      """)
+
     @Test
     fun `list constructor call with generator expression`() = test("""
       expr = list(int(i) for i in '1')
@@ -1761,4 +1769,5 @@ class PyComprehensionAndIteratorTypeTest : PyCodeInsightTestCase() {
       #└ TYPE AsyncGenerator[Literal[42], Any]
       """)
   }
+
 }