From 876cf564581a106ea87c9d2d1c9d3475121d7038 Mon Sep 17 00:00:00 2001
From: Andrey Vokin <andrey.vokin@jetbrains.com>
Date: Wed, 24 Jun 2026 13:21:35 +0200
Subject: [PATCH] PY-89217 PEP-798 - generator expression with unpacking as
 call argument

GitOrigin-RevId: 46422f3e5c7188fcb28f8624ef5f058bd2b10728
---
 .../python/parsing/ExpressionParsing.java     |  9 +++++++-
 .../unpackingInComprehensions.py              |  1 +
 .../testData/psi/UnpackingInComprehensions.py |  1 +
 .../psi/UnpackingInComprehensions.txt         | 23 +++++++++++++++++++
 .../PyComprehensionAndIteratorTypeTest.kt     |  9 ++++++++
 5 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/python/python-parser/src/com/jetbrains/python/parsing/ExpressionParsing.java b/python/python-parser/src/com/jetbrains/python/parsing/ExpressionParsing.java
index aed5c4484a74..a7ec4cb5b2d4 100644
--- a/python/python-parser/src/com/jetbrains/python/parsing/ExpressionParsing.java
+++ b/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)) {
diff --git a/python/testData/inspections/PyCompatibilityInspection/unpackingInComprehensions.py b/python/testData/inspections/PyCompatibilityInspection/unpackingInComprehensions.py
index 49966163ce92..6fe5127ec379 100644
--- a/python/testData/inspections/PyCompatibilityInspection/unpackingInComprehensions.py
+++ b/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)
diff --git a/python/testData/psi/UnpackingInComprehensions.py b/python/testData/psi/UnpackingInComprehensions.py
index 4400c99449df..b976e90b0eeb 100644
--- a/python/testData/psi/UnpackingInComprehensions.py
+++ b/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)
diff --git a/python/testData/psi/UnpackingInComprehensions.txt b/python/testData/psi/UnpackingInComprehensions.txt
index 915a5c80f833..3d00e82e0129 100644
--- a/python/testData/psi/UnpackingInComprehensions.txt
+++ b/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)(')')
diff --git a/python/testSrc/com/jetbrains/python/types/PyComprehensionAndIteratorTypeTest.kt b/python/testSrc/com/jetbrains/python/types/PyComprehensionAndIteratorTypeTest.kt
index ffc0de87f19c..745a443c2ff1 100644
--- a/python/testSrc/com/jetbrains/python/types/PyComprehensionAndIteratorTypeTest.kt
+++ b/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]
       """)
   }
+
 }