python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java - platform/tools/idea - Git at Google

 /*
  * Copyright 2000-2014 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
  * http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 package com.jetbrains.python.psi.impl;

 import com.google.common.collect.Lists;
 import com.intellij.lang.ASTNode;
 import com.intellij.openapi.util.Pair;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiErrorElement;
 import com.intellij.psi.tree.TokenSet;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.IncorrectOperationException;
 import com.intellij.util.SmartList;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.toolbox.FP;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;

 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;

 /**
  * @author yole
  */
 public class PyAssignmentStatementImpl extends PyElementImpl implements PyAssignmentStatement {
   private PyExpression[] myTargets;

   public PyAssignmentStatementImpl(ASTNode astNode) {
     super(astNode);
   }

   @Override
   protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
     pyVisitor.visitPyAssignmentStatement(this);
   }

   public PyExpression[] getTargets() {
     if (myTargets == null) {
       myTargets = calcTargets(false);
     }
     return myTargets;
   }

   @NotNull
   @Override
   public PyExpression[] getRawTargets() {
     return calcTargets(true);
   }

   private PyExpression[] calcTargets(boolean raw) {
     final ASTNode[] eqSigns = getNode().getChildren(TokenSet.create(PyTokenTypes.EQ));
     if (eqSigns.length == 0) {
       return PyExpression.EMPTY_ARRAY;
     }
     final ASTNode lastEq = eqSigns[eqSigns.length - 1];
     List<PyExpression> candidates = new ArrayList<PyExpression>();
     ASTNode node = getNode().getFirstChildNode();
     while (node != null && node != lastEq) {
       final PsiElement psi = node.getPsi();
       if (psi instanceof PyExpression) {
         if (raw) {
           candidates.add((PyExpression) psi);
         }
         else {
           addCandidate(candidates, (PyExpression)psi);
         }
       }
       node = node.getTreeNext();
     }
     List<PyExpression> targets = new ArrayList<PyExpression>();
     for (PyExpression expr : candidates) { // only filter out targets
       if (raw ||
           expr instanceof PyTargetExpression ||
           expr instanceof PyReferenceExpression ||
           expr instanceof PySubscriptionExpression ||
           expr instanceof PySliceExpression) {
         targets.add(expr);
       }
     }
     return targets.toArray(new PyExpression[targets.size()]);
   }

   private static void addCandidate(List<PyExpression> candidates, PyExpression psi) {
     if (psi instanceof PyParenthesizedExpression) {
       addCandidate(candidates, ((PyParenthesizedExpression)psi).getContainedExpression());
     }
     else if (psi instanceof PySequenceExpression) {
       final PyExpression[] pyExpressions = ((PySequenceExpression)psi).getElements();
       for (PyExpression pyExpression : pyExpressions) {
         addCandidate(candidates, pyExpression);
       }
     }
     else if (psi instanceof PyStarExpression) {
       final PyExpression expression = ((PyStarExpression)psi).getExpression();
       if (expression != null) {
         addCandidate(candidates, expression);
       }
     }
     else {
       candidates.add(psi);
     }
   }

   /**
    * @return rightmost expression in statement, which is supposedly the assigned value, or null.
    */
   @Nullable
   public PyExpression getAssignedValue() {
     PsiElement child = getLastChild();
     while (child != null && !(child instanceof PyExpression)) {
       if (child instanceof PsiErrorElement) return null; // incomplete assignment operator can't be analyzed properly, bail out.
       child = child.getPrevSibling();
     }
     return (PyExpression)child;
   }

   @NotNull
   public List<Pair<PyExpression, PyExpression>> getTargetsToValuesMapping() {
     List<Pair<PyExpression, PyExpression>> ret = new SmartList<Pair<PyExpression, PyExpression>>();
     if (!PsiTreeUtil.hasErrorElements(this)) { // no parse errors
       PyExpression[] constituents = PsiTreeUtil.getChildrenOfType(this, PyExpression.class); // "a = b = c" -> [a, b, c]
       if (constituents != null && constituents.length > 1) {
         PyExpression rhs = constituents[constituents.length - 1]; // last
         List<PyExpression> lhses = Lists.newArrayList(constituents);
         if (lhses.size()>0) lhses.remove(lhses.size()-1); // copy all but last; most often it's one element.
         for (PyExpression lhs : lhses) mapToValues(lhs, rhs, ret);
       }
     }
     return ret;
   }

   @Nullable
   public PyExpression getLeftHandSideExpression() {
     PsiElement child = getFirstChild();
     while (child != null && !(child instanceof PyExpression)) {
       if (child instanceof PsiErrorElement) return null; // incomplete assignment operator can't be analyzed properly, bail out.
       child = child.getPrevSibling();
     }
     return (PyExpression)child;
   }

   @Override
   public boolean isAssignmentTo(@NotNull String name) {
     PyExpression lhs = getLeftHandSideExpression();
     return lhs instanceof PyTargetExpression && name.equals(lhs.getName());
   }

   private static void mapToValues(PyExpression lhs, PyExpression rhs, List<Pair<PyExpression, PyExpression>> map) {
     // cast for convenience
     PySequenceExpression lhs_tuple = null;
     PyExpression lhs_one = null;
     if (lhs instanceof PySequenceExpression) lhs_tuple = (PySequenceExpression)lhs;
     else if (lhs != null) lhs_one = lhs;

     PySequenceExpression rhs_tuple = null;
     PyExpression rhs_one = null;
     if (rhs instanceof PyParenthesizedExpression) {
       PyExpression exp = ((PyParenthesizedExpression)rhs).getContainedExpression();
       if (exp instanceof PyTupleExpression)
         rhs_tuple = (PySequenceExpression)exp;
       else
         rhs_one = rhs;
     }
     else if (rhs instanceof PySequenceExpression) rhs_tuple = (PySequenceExpression)rhs;
     else if (rhs != null) rhs_one = rhs;
     //
     if (lhs_one != null) { // single LHS, single RHS (direct mapping) or multiple RHS (packing)
        map.add(Pair.create(lhs_one, rhs));
     }
     else if (lhs_tuple != null && rhs_one != null) { // multiple LHS, single RHS: unpacking
       // PY-2648, PY-2649
       PyElementGenerator elementGenerator = PyElementGenerator.getInstance(rhs_one.getProject());
       final LanguageLevel languageLevel = LanguageLevel.forElement(lhs);
       int counter = 0;
       for (PyExpression tuple_elt : lhs_tuple.getElements()) {
         try {
           final PyExpression expression = elementGenerator.createExpressionFromText(languageLevel, rhs_one.getText() + "[" + counter + "]");
           map.add(Pair.create(tuple_elt, expression));
         }
         catch (IncorrectOperationException e) {
           // not parsed, no problem
         }
         ++counter;
       }
       //  map.addAll(FP.zipList(Arrays.asList(lhs_tuple.getElements()), new RepeatIterable<PyExpression>(rhs_one)));
     }
     else if (lhs_tuple != null && rhs_tuple != null) { // multiple both sides: piecewise mapping
       map.addAll(FP.zipList(Arrays.asList(lhs_tuple.getElements()), Arrays.asList(rhs_tuple.getElements()), null, null));
     }
   }

   @NotNull
   public Iterable<PyElement> iterateNames() {
     final List<PyExpression> expressions = PyUtil.flattenedParensAndStars(getTargets());
     List<PyElement> result = new ArrayList<PyElement>();
     for (PyExpression expression : expressions) {
       if (expression instanceof PyQualifiedExpression && ((PyQualifiedExpression)expression).isQualified()) {
         continue;
       }
       result.add(expression);
     }
     return result;

   }

   public PyElement getElementNamed(final String the_name) {
     // performance: check simple case first
     PyExpression[] targets = getTargets();
     if (targets.length == 1 && targets[0] instanceof PyTargetExpression) {
       PyTargetExpression target = (PyTargetExpression)targets[0];
       return !target.isQualified() && the_name.equals(target.getName()) ? target : null;
     }
     return IterHelper.findName(iterateNames(), the_name);
   }

   public boolean mustResolveOutside() {
     return true; // a = a+1 resolves 'a' outside itself.
   }

   @Override
   public void subtreeChanged() {
     super.subtreeChanged();
     myTargets = null;
   }
 }
	/*
	* Copyright 2000-2014 JetBrains s.r.o.
	*
	* Licensed under the Apache License, Version 2.0 (the "License");
	* you may not use this file except in compliance with the License.
	* You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/
	package com.jetbrains.python.psi.impl;

	import com.google.common.collect.Lists;
	import com.intellij.lang.ASTNode;
	import com.intellij.openapi.util.Pair;
	import com.intellij.psi.PsiElement;
	import com.intellij.psi.PsiErrorElement;
	import com.intellij.psi.tree.TokenSet;
	import com.intellij.psi.util.PsiTreeUtil;
	import com.intellij.util.IncorrectOperationException;
	import com.intellij.util.SmartList;
	import com.jetbrains.python.PyTokenTypes;
	import com.jetbrains.python.psi.*;
	import com.jetbrains.python.toolbox.FP;
	import org.jetbrains.annotations.NotNull;
	import org.jetbrains.annotations.Nullable;

	import java.util.ArrayList;
	import java.util.Arrays;
	import java.util.List;

	/**
	* @author yole
	*/
	public class PyAssignmentStatementImpl extends PyElementImpl implements PyAssignmentStatement {
	private PyExpression[] myTargets;

	public PyAssignmentStatementImpl(ASTNode astNode) {
	super(astNode);
	}

	@Override
	protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
	pyVisitor.visitPyAssignmentStatement(this);
	}

	public PyExpression[] getTargets() {
	if (myTargets == null) {
	myTargets = calcTargets(false);
	}
	return myTargets;
	}

	@NotNull
	@Override
	public PyExpression[] getRawTargets() {
	return calcTargets(true);
	}

	private PyExpression[] calcTargets(boolean raw) {
	final ASTNode[] eqSigns = getNode().getChildren(TokenSet.create(PyTokenTypes.EQ));
	if (eqSigns.length == 0) {
	return PyExpression.EMPTY_ARRAY;
	}
	final ASTNode lastEq = eqSigns[eqSigns.length - 1];
	List<PyExpression> candidates = new ArrayList<PyExpression>();
	ASTNode node = getNode().getFirstChildNode();
	while (node != null && node != lastEq) {
	final PsiElement psi = node.getPsi();
	if (psi instanceof PyExpression) {
	if (raw) {
	candidates.add((PyExpression) psi);
	}
	else {
	addCandidate(candidates, (PyExpression)psi);
	}
	}
	node = node.getTreeNext();
	}
	List<PyExpression> targets = new ArrayList<PyExpression>();
	for (PyExpression expr : candidates) { // only filter out targets
	if (raw \|\|
	expr instanceof PyTargetExpression \|\|
	expr instanceof PyReferenceExpression \|\|
	expr instanceof PySubscriptionExpression \|\|
	expr instanceof PySliceExpression) {
	targets.add(expr);
	}
	}
	return targets.toArray(new PyExpression[targets.size()]);
	}

	private static void addCandidate(List<PyExpression> candidates, PyExpression psi) {
	if (psi instanceof PyParenthesizedExpression) {
	addCandidate(candidates, ((PyParenthesizedExpression)psi).getContainedExpression());
	}
	else if (psi instanceof PySequenceExpression) {
	final PyExpression[] pyExpressions = ((PySequenceExpression)psi).getElements();
	for (PyExpression pyExpression : pyExpressions) {
	addCandidate(candidates, pyExpression);
	}
	}
	else if (psi instanceof PyStarExpression) {
	final PyExpression expression = ((PyStarExpression)psi).getExpression();
	if (expression != null) {
	addCandidate(candidates, expression);
	}
	}
	else {
	candidates.add(psi);
	}
	}

	/**
	* @return rightmost expression in statement, which is supposedly the assigned value, or null.
	*/
	@Nullable
	public PyExpression getAssignedValue() {
	PsiElement child = getLastChild();
	while (child != null && !(child instanceof PyExpression)) {
	if (child instanceof PsiErrorElement) return null; // incomplete assignment operator can't be analyzed properly, bail out.
	child = child.getPrevSibling();
	}
	return (PyExpression)child;
	}

	@NotNull
	public List<Pair<PyExpression, PyExpression>> getTargetsToValuesMapping() {
	List<Pair<PyExpression, PyExpression>> ret = new SmartList<Pair<PyExpression, PyExpression>>();
	if (!PsiTreeUtil.hasErrorElements(this)) { // no parse errors
	PyExpression[] constituents = PsiTreeUtil.getChildrenOfType(this, PyExpression.class); // "a = b = c" -> [a, b, c]
	if (constituents != null && constituents.length > 1) {
	PyExpression rhs = constituents[constituents.length - 1]; // last
	List<PyExpression> lhses = Lists.newArrayList(constituents);
	if (lhses.size()>0) lhses.remove(lhses.size()-1); // copy all but last; most often it's one element.
	for (PyExpression lhs : lhses) mapToValues(lhs, rhs, ret);
	}
	}
	return ret;
	}

	@Nullable
	public PyExpression getLeftHandSideExpression() {
	PsiElement child = getFirstChild();
	while (child != null && !(child instanceof PyExpression)) {
	if (child instanceof PsiErrorElement) return null; // incomplete assignment operator can't be analyzed properly, bail out.
	child = child.getPrevSibling();
	}
	return (PyExpression)child;
	}

	@Override
	public boolean isAssignmentTo(@NotNull String name) {
	PyExpression lhs = getLeftHandSideExpression();
	return lhs instanceof PyTargetExpression && name.equals(lhs.getName());
	}

	private static void mapToValues(PyExpression lhs, PyExpression rhs, List<Pair<PyExpression, PyExpression>> map) {
	// cast for convenience
	PySequenceExpression lhs_tuple = null;
	PyExpression lhs_one = null;
	if (lhs instanceof PySequenceExpression) lhs_tuple = (PySequenceExpression)lhs;
	else if (lhs != null) lhs_one = lhs;

	PySequenceExpression rhs_tuple = null;
	PyExpression rhs_one = null;
	if (rhs instanceof PyParenthesizedExpression) {
	PyExpression exp = ((PyParenthesizedExpression)rhs).getContainedExpression();
	if (exp instanceof PyTupleExpression)
	rhs_tuple = (PySequenceExpression)exp;
	else
	rhs_one = rhs;
	}
	else if (rhs instanceof PySequenceExpression) rhs_tuple = (PySequenceExpression)rhs;
	else if (rhs != null) rhs_one = rhs;
	//
	if (lhs_one != null) { // single LHS, single RHS (direct mapping) or multiple RHS (packing)
	map.add(Pair.create(lhs_one, rhs));
	}
	else if (lhs_tuple != null && rhs_one != null) { // multiple LHS, single RHS: unpacking
	// PY-2648, PY-2649
	PyElementGenerator elementGenerator = PyElementGenerator.getInstance(rhs_one.getProject());
	final LanguageLevel languageLevel = LanguageLevel.forElement(lhs);
	int counter = 0;
	for (PyExpression tuple_elt : lhs_tuple.getElements()) {
	try {
	final PyExpression expression = elementGenerator.createExpressionFromText(languageLevel, rhs_one.getText() + "[" + counter + "]");
	map.add(Pair.create(tuple_elt, expression));
	}
	catch (IncorrectOperationException e) {
	// not parsed, no problem
	}
	++counter;
	}
	// map.addAll(FP.zipList(Arrays.asList(lhs_tuple.getElements()), new RepeatIterable<PyExpression>(rhs_one)));
	}
	else if (lhs_tuple != null && rhs_tuple != null) { // multiple both sides: piecewise mapping
	map.addAll(FP.zipList(Arrays.asList(lhs_tuple.getElements()), Arrays.asList(rhs_tuple.getElements()), null, null));
	}
	}

	@NotNull
	public Iterable<PyElement> iterateNames() {
	final List<PyExpression> expressions = PyUtil.flattenedParensAndStars(getTargets());
	List<PyElement> result = new ArrayList<PyElement>();
	for (PyExpression expression : expressions) {
	if (expression instanceof PyQualifiedExpression && ((PyQualifiedExpression)expression).isQualified()) {
	continue;
	}
	result.add(expression);
	}
	return result;

	}

	public PyElement getElementNamed(final String the_name) {
	// performance: check simple case first
	PyExpression[] targets = getTargets();
	if (targets.length == 1 && targets[0] instanceof PyTargetExpression) {
	PyTargetExpression target = (PyTargetExpression)targets[0];
	return !target.isQualified() && the_name.equals(target.getName()) ? target : null;
	}
	return IterHelper.findName(iterateNames(), the_name);
	}

	public boolean mustResolveOutside() {
	return true; // a = a+1 resolves 'a' outside itself.
	}

	@Override
	public void subtreeChanged() {
	super.subtreeChanged();
	myTargets = null;
	}
	}