From 6fc5759c3382eb90ed16efbbd5fe2fab732e765c Mon Sep 17 00:00:00 2001 From: Dmitry Jemerov Date: Tue, 17 Aug 2010 14:39:00 +0400 Subject: [PATCH] unfold star expressions when iterating names in 'for' and similar statements (PY-1525) --- .../src/com/jetbrains/python/psi/PyUtil.java | 34 ++++++++++--------- .../psi/impl/PyAssignmentStatementImpl.java | 4 +-- .../psi/impl/PyComprehensionElementImpl.java | 5 ++- .../python/psi/impl/PyExceptPartImpl.java | 28 +++++++++------ .../python/psi/impl/PyForStatementImpl.java | 6 ++-- .../psi/impl/PyGeneratorExpressionImpl.java | 6 ++-- .../testData/resolve/StarUnpackingInLoop.py | 5 +++ .../com/jetbrains/python/PyResolveTest.java | 13 +++---- .../python/fixtures/PyResolveTestCase.java | 12 +++++++ 9 files changed, 65 insertions(+), 48 deletions(-) create mode 100644 python/testData/resolve/StarUnpackingInLoop.py diff --git a/python/src/com/jetbrains/python/psi/PyUtil.java b/python/src/com/jetbrains/python/psi/PyUtil.java index 19edb335b226..b35c7e49e9aa 100644 --- a/python/src/com/jetbrains/python/psi/PyUtil.java +++ b/python/src/com/jetbrains/python/psi/PyUtil.java @@ -110,27 +110,24 @@ public class PyUtil { /** * @see PyUtil#flattenedParens */ - protected static List _unfoldParenExprs(T[] targets, List receiver, boolean unfoldListLiterals) { + protected static List _unfoldParenExprs(PyExpression[] targets, List receiver, + boolean unfoldListLiterals, boolean unfoldStarExpressions) { // NOTE: this proliferation of instanceofs is not very beautiful. Maybe rewrite using a visitor. - for (T exp : targets) { + for (PyExpression exp : targets) { if (exp instanceof PyParenthesizedExpression) { final PyParenthesizedExpression parex = (PyParenthesizedExpression)exp; - PyExpression cont = parex.getContainedExpression(); - if (cont instanceof PyTupleExpression) { - final PyTupleExpression tupex = (PyTupleExpression)cont; - _unfoldParenExprs((T[])tupex.getElements(), receiver, unfoldListLiterals); - } - else { - receiver.add(exp); - } + _unfoldParenExprs(new PyExpression[] { parex.getContainedExpression() }, receiver, unfoldListLiterals, unfoldStarExpressions); } else if (exp instanceof PyTupleExpression) { final PyTupleExpression tupex = (PyTupleExpression)exp; - _unfoldParenExprs((T[])tupex.getElements(), receiver, unfoldListLiterals); + _unfoldParenExprs(tupex.getElements(), receiver, unfoldListLiterals, unfoldStarExpressions); } else if (exp instanceof PyListLiteralExpression && unfoldListLiterals) { final PyListLiteralExpression listLiteral = (PyListLiteralExpression) exp; - _unfoldParenExprs((T[]) listLiteral.getElements(), receiver, unfoldListLiterals); + _unfoldParenExprs(listLiteral.getElements(), receiver, unfoldListLiterals, unfoldStarExpressions); + } + else if (exp instanceof PyStarExpression && unfoldStarExpressions) { + _unfoldParenExprs(new PyExpression[] { ((PyStarExpression) exp).getExpression() }, receiver, unfoldListLiterals, unfoldStarExpressions); } else { receiver.add(exp); @@ -150,13 +147,18 @@ public class PyUtil { * @return the list of flattened expressions. */ @NotNull - public static List flattenedParens(T... targets) { - return _unfoldParenExprs(targets, new ArrayList(targets.length), false); + public static List flattenedParens(PyExpression... targets) { + return _unfoldParenExprs(targets, new ArrayList(targets.length), false, false); } @NotNull - public static List flattenedParensAndLists(T... targets) { - return _unfoldParenExprs(targets, new ArrayList(targets.length), true); + public static List flattenedParensAndLists(PyExpression... targets) { + return _unfoldParenExprs(targets, new ArrayList(targets.length), true, true); + } + + @NotNull + public static List flattenedParensAndStars(PyExpression... targets) { + return _unfoldParenExprs(targets, new ArrayList(targets.length), false, true); } /** diff --git a/python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java b/python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java index 5b186e57d92c..4bffb53f35a4 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java @@ -153,9 +153,7 @@ public class PyAssignmentStatementImpl extends PyElementImpl implements PyAssign @NotNull public Iterable iterateNames() { - PyElement[] targets = getTargets(); - // return _unfoldParenExprs(targets, new ArrayList(targets.length)); - return PyUtil.flattenedParens(targets); + return new ArrayList(PyUtil.flattenedParensAndStars(getTargets())); } public PyElement getElementNamed(final String the_name) { diff --git a/python/src/com/jetbrains/python/psi/impl/PyComprehensionElementImpl.java b/python/src/com/jetbrains/python/psi/impl/PyComprehensionElementImpl.java index 784924b58b15..fb4f854cd346 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyComprehensionElementImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyComprehensionElementImpl.java @@ -101,14 +101,13 @@ public abstract class PyComprehensionElementImpl extends PyElementImpl implement public Iterable iterateNames() { // extract whatever names are defined in "for" components List fors = getForComponents(); - PyElement[] for_targets = new PyElement[fors.size()]; + PyExpression[] for_targets = new PyExpression[fors.size()]; int i = 0; for (ComprhForComponent for_comp : fors) { for_targets[i] = for_comp.getIteratorVariable(); i += 1; } - List name_refs = PyUtil.flattenedParensAndLists(for_targets); - return name_refs; + return new ArrayList(PyUtil.flattenedParensAndLists(for_targets)); } public PsiElement getElementNamed(final String the_name) { diff --git a/python/src/com/jetbrains/python/psi/impl/PyExceptPartImpl.java b/python/src/com/jetbrains/python/psi/impl/PyExceptPartImpl.java index 09796bb2247f..65139c551819 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyExceptPartImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyExceptPartImpl.java @@ -6,33 +6,39 @@ import com.jetbrains.python.psi.*; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; +import java.util.ArrayList; + /** * @author dcheryasov */ public class PyExceptPartImpl extends PyElementImpl implements PyExceptPart { public PyExceptPartImpl(ASTNode astNode) { - super(astNode); + super(astNode); } - @Override protected void acceptPyVisitor(PyElementVisitor pyVisitor) { - pyVisitor.visitPyExceptBlock(this); + @Override + protected void acceptPyVisitor(PyElementVisitor pyVisitor) { + pyVisitor.visitPyExceptBlock(this); } - public @Nullable PyExpression getExceptClass() { - return childToPsi(PyElementTypes.EXPRESSIONS, 0); + @Nullable + public PyExpression getExceptClass() { + return childToPsi(PyElementTypes.EXPRESSIONS, 0); } - public @Nullable PyExpression getTarget() { - return childToPsi(PyElementTypes.EXPRESSIONS, 1); + @Nullable + public PyExpression getTarget() { + return childToPsi(PyElementTypes.EXPRESSIONS, 1); } - public @NotNull PyStatementList getStatementList() { - return childToPsiNotNull(PyElementTypes.STATEMENT_LIST); + @NotNull + public PyStatementList getStatementList() { + return childToPsiNotNull(PyElementTypes.STATEMENT_LIST); } @NotNull public Iterable iterateNames() { - return PyUtil.flattenedParens(getTarget()); + return new ArrayList(PyUtil.flattenedParensAndStars(getTarget())); } public PyElement getElementNamed(final String the_name) { @@ -40,6 +46,6 @@ public class PyExceptPartImpl extends PyElementImpl implements PyExceptPart { } public boolean mustResolveOutside() { - return false; + return false; } } diff --git a/python/src/com/jetbrains/python/psi/impl/PyForStatementImpl.java b/python/src/com/jetbrains/python/psi/impl/PyForStatementImpl.java index a9d29547dc93..1f13bfa64221 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyForStatementImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyForStatementImpl.java @@ -1,13 +1,11 @@ package com.jetbrains.python.psi.impl; import com.intellij.lang.ASTNode; -import com.intellij.psi.PsiElement; -import com.intellij.psi.ResolveState; -import com.intellij.psi.scope.PsiScopeProcessor; import com.jetbrains.python.PyElementTypes; import com.jetbrains.python.psi.*; import org.jetbrains.annotations.NotNull; +import java.util.ArrayList; import java.util.Collections; public class PyForStatementImpl extends PyPartitionedElementImpl implements PyForStatement { @@ -34,7 +32,7 @@ public class PyForStatementImpl extends PyPartitionedElementImpl implements PyFo PyExpression tgt = getForPart().getTarget(); if (tgt instanceof PyReferenceExpression) return Collections.singleton(tgt); else { - return PyUtil.flattenedParens(new PyElement[]{tgt}); + return new ArrayList(PyUtil.flattenedParensAndStars(tgt)); } } diff --git a/python/src/com/jetbrains/python/psi/impl/PyGeneratorExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyGeneratorExpressionImpl.java index c28391c928d1..bf8bb1641d11 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyGeneratorExpressionImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyGeneratorExpressionImpl.java @@ -7,6 +7,7 @@ import com.jetbrains.python.psi.types.PyType; import com.jetbrains.python.psi.types.TypeEvalContext; import org.jetbrains.annotations.NotNull; +import java.util.ArrayList; import java.util.List; /** @@ -30,14 +31,13 @@ public class PyGeneratorExpressionImpl extends PyComprehensionElementImpl implem public Iterable iterateNames() { // extract whatever names are defined in "for" components List fors = getForComponents(); - PyElement[] for_targets = new PyElement[fors.size()]; + PyExpression[] for_targets = new PyExpression[fors.size()]; int i = 0; for (ComprhForComponent for_comp : fors) { for_targets[i] = for_comp.getIteratorVariable(); i += 1; } - List name_refs = PyUtil.flattenedParens(for_targets); - return name_refs; + return new ArrayList(PyUtil.flattenedParensAndStars(for_targets)); } public PsiElement getElementNamed(final String the_name) { diff --git a/python/testData/resolve/StarUnpackingInLoop.py b/python/testData/resolve/StarUnpackingInLoop.py new file mode 100644 index 000000000000..6c28f6b7a214 --- /dev/null +++ b/python/testData/resolve/StarUnpackingInLoop.py @@ -0,0 +1,5 @@ +def implicit_assignment(): + seq = [(1, 2, 3), (4, 5, 6, 7)] + for a, *bbb in seq: + print(bbb) +# \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyResolveTest.java b/python/testSrc/com/jetbrains/python/PyResolveTest.java index 7f2a33f9b70e..c409bb4e265a 100644 --- a/python/testSrc/com/jetbrains/python/PyResolveTest.java +++ b/python/testSrc/com/jetbrains/python/PyResolveTest.java @@ -361,13 +361,10 @@ public class PyResolveTest extends PyResolveTestCase { } public void testStarUnpacking() { // PY-1459 - PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), LanguageLevel.PYTHON30); - try { - final PsiElement element = doResolve(); - assertInstanceOf(element, PyTargetExpression.class); - } - finally { - PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), null); - } + assertResolvesTo(LanguageLevel.PYTHON30, PyTargetExpression.class, "heads"); + } + + public void testStarUnpackingInLoop() { // PY-1525 + assertResolvesTo(LanguageLevel.PYTHON30, PyTargetExpression.class, "bbb"); } } \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/fixtures/PyResolveTestCase.java b/python/testSrc/com/jetbrains/python/fixtures/PyResolveTestCase.java index 4675856bf2b5..a4d7babf66a6 100644 --- a/python/testSrc/com/jetbrains/python/fixtures/PyResolveTestCase.java +++ b/python/testSrc/com/jetbrains/python/fixtures/PyResolveTestCase.java @@ -8,6 +8,8 @@ import com.intellij.openapi.vfs.VfsUtil; import com.intellij.openapi.vfs.VirtualFile; import com.intellij.psi.*; import com.intellij.testFramework.TestDataFile; +import com.jetbrains.python.psi.LanguageLevel; +import com.jetbrains.python.psi.impl.PythonLanguageLevelPusher; import org.jetbrains.annotations.NonNls; import java.io.File; @@ -47,6 +49,16 @@ public abstract class PyResolveTestCase extends PyLightFixtureTestCase { protected abstract PsiElement doResolve() throws Exception; + protected T assertResolvesTo(final LanguageLevel langLevel, final Class aClass, final String name) { + PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), langLevel); + try { + return assertResolvesTo(aClass, name, null); + } + finally { + PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), null); + } + } + protected T assertResolvesTo(final Class aClass, final String name) { return assertResolvesTo(aClass, name, null); }