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