com.jetbrains.python.psi.impl.PyAssignmentStatementImpl Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of python-community Show documentation
A packaging of the IntelliJ Community Edition python-community library. This is release number 1 of trunk branch 142.
The newest version!
/*
 * 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 candidates = new ArrayList();
    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 targets = new ArrayList();
    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 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> getTargetsToValuesMapping() {
    List> ret = new SmartList>();
    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 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> 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(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 iterateNames() {
    final List expressions = PyUtil.flattenedParensAndStars(getTargets());
    List result = new ArrayList();
    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;
  }
}