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 @@
{*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.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]
""")
}
+
}