diff --git a/python/src/com/jetbrains/python/validation/StarAnnotator.java b/python/src/com/jetbrains/python/validation/StarAnnotator.java
index 55cda35fadc9..73cb0eb06d99 100644
--- a/python/src/com/jetbrains/python/validation/StarAnnotator.java
+++ b/python/src/com/jetbrains/python/validation/StarAnnotator.java
@@ -15,8 +15,13 @@
*/
package com.jetbrains.python.validation;
+import com.intellij.psi.PsiElement;
import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyParameterTypeList;
+import com.jetbrains.python.psi.PyReturnStatement;
import com.jetbrains.python.psi.PyStarExpression;
+import com.jetbrains.python.psi.PyTupleExpression;
+import com.jetbrains.python.psi.PyYieldExpression;
+import org.jetbrains.annotations.NotNull;
/**
* @author yole
@@ -25,8 +30,21 @@ public class StarAnnotator extends PyAnnotator {
@Override
public void visitPyStarExpression(PyStarExpression node) {
super.visitPyStarExpression(node);
- if (!node.isAssignmentTarget() && !node.isUnpacking() && !(node.getParent() instanceof PyParameterTypeList)) {
+ if (!node.isAssignmentTarget() && !allowedUnpacking(node) && !(node.getParent() instanceof PyParameterTypeList)) {
getHolder().createErrorAnnotation(node, "Can't use starred expression here");
}
}
+
+ private static boolean allowedUnpacking(@NotNull PyStarExpression starExpression) {
+ if (!starExpression.isUnpacking()) {
+ return false;
+ }
+
+ final PsiElement parent = starExpression.getParent();
+ if (parent instanceof PyTupleExpression && (parent.getParent() instanceof PyReturnStatement ||
+ parent.getParent() instanceof PyYieldExpression)) {
+ return false;
+ }
+ return true;
+ }
}
diff --git a/python/testData/highlighting/unpackingStar.py b/python/testData/highlighting/unpackingStar.py
index 59d21937a8ff..5fa38f239683 100644
--- a/python/testData/highlighting/unpackingStar.py
+++ b/python/testData/highlighting/unpackingStar.py
@@ -7,4 +7,13 @@ if *x:
pass
1 + (*x)
-1 + (*x,)
\ No newline at end of file
+1 + (*x,)
+
+
+def f(x):
+ return x, *x
+
+
+def g(x):
+ yield from x, *x
+ yield x, *x
\ No newline at end of file