unfold star expressions when iterating names in 'for' and similar statements (PY-1525)

This commit is contained in:
Dmitry Jemerov
2010-08-17 21:35:42 +04:00
parent e446210a0d
commit 6fc5759c33
9 changed files with 65 additions and 48 deletions
+18 -16
View File
@@ -110,27 +110,24 @@ public class PyUtil {
/**
* @see PyUtil#flattenedParens
*/
protected static <T extends PyElement> List<T> _unfoldParenExprs(T[] targets, List<T> receiver, boolean unfoldListLiterals) {
protected static List<PyExpression> _unfoldParenExprs(PyExpression[] targets, List<PyExpression> 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 <T extends PyElement> List<T> flattenedParens(T... targets) {
return _unfoldParenExprs(targets, new ArrayList<T>(targets.length), false);
public static List<PyExpression> flattenedParens(PyExpression... targets) {
return _unfoldParenExprs(targets, new ArrayList<PyExpression>(targets.length), false, false);
}
@NotNull
public static <T extends PyElement> List<T> flattenedParensAndLists(T... targets) {
return _unfoldParenExprs(targets, new ArrayList<T>(targets.length), true);
public static List<PyExpression> flattenedParensAndLists(PyExpression... targets) {
return _unfoldParenExprs(targets, new ArrayList<PyExpression>(targets.length), true, true);
}
@NotNull
public static List<PyExpression> flattenedParensAndStars(PyExpression... targets) {
return _unfoldParenExprs(targets, new ArrayList<PyExpression>(targets.length), false, true);
}
/**
@@ -153,9 +153,7 @@ public class PyAssignmentStatementImpl extends PyElementImpl implements PyAssign
@NotNull
public Iterable<PyElement> iterateNames() {
PyElement[] targets = getTargets();
// return _unfoldParenExprs(targets, new ArrayList<PyElement>(targets.length));
return PyUtil.flattenedParens(targets);
return new ArrayList<PyElement>(PyUtil.flattenedParensAndStars(getTargets()));
}
public PyElement getElementNamed(final String the_name) {
@@ -101,14 +101,13 @@ public abstract class PyComprehensionElementImpl extends PyElementImpl implement
public Iterable<PyElement> iterateNames() {
// extract whatever names are defined in "for" components
List<ComprhForComponent> 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<PyElement> name_refs = PyUtil.flattenedParensAndLists(for_targets);
return name_refs;
return new ArrayList<PyElement>(PyUtil.flattenedParensAndLists(for_targets));
}
public PsiElement getElementNamed(final String the_name) {
@@ -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<PyElement> iterateNames() {
return PyUtil.<PyElement>flattenedParens(getTarget());
return new ArrayList<PyElement>(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;
}
}
@@ -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.<PyElement>singleton(tgt);
else {
return PyUtil.flattenedParens(new PyElement[]{tgt});
return new ArrayList<PyElement>(PyUtil.flattenedParensAndStars(tgt));
}
}
@@ -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<PyElement> iterateNames() {
// extract whatever names are defined in "for" components
List<ComprhForComponent> 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<PyElement> name_refs = PyUtil.flattenedParens(for_targets);
return name_refs;
return new ArrayList<PyElement>(PyUtil.flattenedParensAndStars(for_targets));
}
public PsiElement getElementNamed(final String the_name) {
@@ -0,0 +1,5 @@
def implicit_assignment():
seq = [(1, 2, 3), (4, 5, 6, 7)]
for a, *bbb in seq:
print(bbb)
# <ref>
@@ -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");
}
}
@@ -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 extends PsiElement> T assertResolvesTo(final LanguageLevel langLevel, final Class<T> aClass, final String name) {
PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), langLevel);
try {
return assertResolvesTo(aClass, name, null);
}
finally {
PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), null);
}
}
protected <T extends PsiElement> T assertResolvesTo(final Class<T> aClass, final String name) {
return assertResolvesTo(aClass, name, null);
}