PY-20744 Parse PEP-526 variable annotations
[idea/community.git] / python / src / com / jetbrains / python / psi / impl / PyAssignmentStatementImpl.java
1 /*
2  * Copyright 2000-2014 JetBrains s.r.o.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package com.jetbrains.python.psi.impl;
17
18 import com.google.common.collect.Lists;
19 import com.intellij.lang.ASTNode;
20 import com.intellij.openapi.util.Pair;
21 import com.intellij.psi.PsiElement;
22 import com.intellij.psi.PsiErrorElement;
23 import com.intellij.psi.PsiNamedElement;
24 import com.intellij.psi.tree.TokenSet;
25 import com.intellij.psi.util.PsiTreeUtil;
26 import com.intellij.util.IncorrectOperationException;
27 import com.intellij.util.SmartList;
28 import com.jetbrains.python.PyTokenTypes;
29 import com.jetbrains.python.psi.*;
30 import com.jetbrains.python.toolbox.FP;
31 import org.jetbrains.annotations.NotNull;
32 import org.jetbrains.annotations.Nullable;
33
34 import java.util.ArrayList;
35 import java.util.Arrays;
36 import java.util.List;
37
38 /**
39  * @author yole
40  */
41 public class PyAssignmentStatementImpl extends PyElementImpl implements PyAssignmentStatement {
42   private PyExpression[] myTargets;
43
44   public PyAssignmentStatementImpl(ASTNode astNode) {
45     super(astNode);
46   }
47
48   @Override
49   protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
50     pyVisitor.visitPyAssignmentStatement(this);
51   }
52
53   public PyExpression[] getTargets() {
54     if (myTargets == null) {
55       myTargets = calcTargets(false);
56     }
57     return myTargets;
58   }
59
60   @NotNull
61   @Override
62   public PyExpression[] getRawTargets() {
63     return calcTargets(true);
64   }
65
66   private PyExpression[] calcTargets(boolean raw) {
67     final ASTNode[] eqSigns = getNode().getChildren(TokenSet.create(PyTokenTypes.EQ));
68     if (eqSigns.length == 0) {
69       return PyExpression.EMPTY_ARRAY;
70     }
71     final ASTNode lastEq = eqSigns[eqSigns.length - 1];
72     List<PyExpression> candidates = new ArrayList<>();
73     ASTNode node = getNode().getFirstChildNode();
74     while (node != null && node != lastEq) {
75       final PsiElement psi = node.getPsi();
76       if (psi instanceof PyExpression) {
77         if (raw) {
78           candidates.add((PyExpression) psi);
79         }
80         else {
81           addCandidate(candidates, (PyExpression)psi);
82         }
83       }
84       node = node.getTreeNext();
85     }
86     List<PyExpression> targets = new ArrayList<>();
87     for (PyExpression expr : candidates) { // only filter out targets
88       if (raw ||
89           expr instanceof PyTargetExpression ||
90           expr instanceof PyReferenceExpression ||
91           expr instanceof PySubscriptionExpression ||
92           expr instanceof PySliceExpression) {
93         targets.add(expr);
94       }
95     }
96     return targets.toArray(new PyExpression[targets.size()]);
97   }
98
99   @Nullable
100   @Override
101   public PyAnnotation getAnnotation() {
102     return findChildByClass(PyAnnotation.class);
103   }
104
105   private static void addCandidate(List<PyExpression> candidates, PyExpression psi) {
106     if (psi instanceof PyParenthesizedExpression) {
107       addCandidate(candidates, ((PyParenthesizedExpression)psi).getContainedExpression());
108     }
109     else if (psi instanceof PySequenceExpression) {
110       final PyExpression[] pyExpressions = ((PySequenceExpression)psi).getElements();
111       for (PyExpression pyExpression : pyExpressions) {
112         addCandidate(candidates, pyExpression);
113       }
114     }
115     else if (psi instanceof PyStarExpression) {
116       final PyExpression expression = ((PyStarExpression)psi).getExpression();
117       if (expression != null) {
118         addCandidate(candidates, expression);
119       }
120     }
121     else {
122       candidates.add(psi);
123     }
124   }
125
126   /**
127    * @return rightmost expression in statement, which is supposedly the assigned value, or null.
128    */
129   @Nullable
130   public PyExpression getAssignedValue() {
131     PsiElement child = getLastChild();
132     while (child != null && !(child instanceof PyExpression)) {
133       if (child instanceof PsiErrorElement) return null; // incomplete assignment operator can't be analyzed properly, bail out.
134       child = child.getPrevSibling();
135     }
136     return (PyExpression)child;
137   }
138
139   @NotNull
140   public List<Pair<PyExpression, PyExpression>> getTargetsToValuesMapping() {
141     List<Pair<PyExpression, PyExpression>> ret = new SmartList<>();
142     if (!PsiTreeUtil.hasErrorElements(this)) { // no parse errors
143       PyExpression[] constituents = PsiTreeUtil.getChildrenOfType(this, PyExpression.class); // "a = b = c" -> [a, b, c]
144       if (constituents != null && constituents.length > 1) {
145         PyExpression rhs = constituents[constituents.length - 1]; // last
146         List<PyExpression> lhses = Lists.newArrayList(constituents);
147         if (lhses.size()>0) lhses.remove(lhses.size()-1); // copy all but last; most often it's one element.
148         for (PyExpression lhs : lhses) mapToValues(lhs, rhs, ret);
149       }
150     }
151     return ret;
152   }
153
154   @Nullable
155   public PyExpression getLeftHandSideExpression() {
156     PsiElement child = getFirstChild();
157     while (child != null && !(child instanceof PyExpression)) {
158       if (child instanceof PsiErrorElement) return null; // incomplete assignment operator can't be analyzed properly, bail out.
159       child = child.getPrevSibling();
160     }
161     return (PyExpression)child;
162   }
163
164   @Override
165   public boolean isAssignmentTo(@NotNull String name) {
166     PyExpression lhs = getLeftHandSideExpression();
167     return lhs instanceof PyTargetExpression && name.equals(lhs.getName());
168   }
169
170   private static void mapToValues(PyExpression lhs, PyExpression rhs, List<Pair<PyExpression, PyExpression>> map) {
171     // cast for convenience
172     PySequenceExpression lhs_tuple = null;
173     PyExpression lhs_one = null;
174     if (lhs instanceof PySequenceExpression) lhs_tuple = (PySequenceExpression)lhs;
175     else if (lhs != null) lhs_one = lhs;
176     
177     PySequenceExpression rhs_tuple = null;
178     PyExpression rhs_one = null;
179     if (rhs instanceof PyParenthesizedExpression) {
180       PyExpression exp = ((PyParenthesizedExpression)rhs).getContainedExpression();
181       if (exp instanceof PyTupleExpression)
182         rhs_tuple = (PySequenceExpression)exp;
183       else
184         rhs_one = rhs;
185     }
186     else if (rhs instanceof PySequenceExpression) rhs_tuple = (PySequenceExpression)rhs;
187     else if (rhs != null) rhs_one = rhs;
188     //
189     if (lhs_one != null) { // single LHS, single RHS (direct mapping) or multiple RHS (packing)
190        map.add(Pair.create(lhs_one, rhs));
191     }
192     else if (lhs_tuple != null && rhs_one != null) { // multiple LHS, single RHS: unpacking
193       // PY-2648, PY-2649
194       PyElementGenerator elementGenerator = PyElementGenerator.getInstance(rhs_one.getProject());
195       final LanguageLevel languageLevel = LanguageLevel.forElement(lhs);
196       int counter = 0;
197       for (PyExpression tuple_elt : lhs_tuple.getElements()) {
198         try {
199           final PyExpression expression = elementGenerator.createExpressionFromText(languageLevel, rhs_one.getText() + "[" + counter + "]");
200           map.add(Pair.create(tuple_elt, expression));
201         }
202         catch (IncorrectOperationException e) {
203           // not parsed, no problem
204         }
205         ++counter;
206       }
207       //  map.addAll(FP.zipList(Arrays.asList(lhs_tuple.getElements()), new RepeatIterable<PyExpression>(rhs_one)));
208     }
209     else if (lhs_tuple != null && rhs_tuple != null) { // multiple both sides: piecewise mapping
210       map.addAll(FP.zipList(Arrays.asList(lhs_tuple.getElements()), Arrays.asList(rhs_tuple.getElements()), null, null));
211     }
212   }
213
214   @NotNull
215   public List<PsiNamedElement> getNamedElements() {
216     final List<PyExpression> expressions = PyUtil.flattenedParensAndStars(getTargets());
217     List<PsiNamedElement> result = new ArrayList<>();
218     for (PyExpression expression : expressions) {
219       if (expression instanceof PyQualifiedExpression && ((PyQualifiedExpression)expression).isQualified()) {
220         continue;
221       }
222       if (expression instanceof PsiNamedElement) {
223         result.add((PsiNamedElement)expression);
224       }
225     }
226     return result;
227
228   }
229
230   @Nullable
231   public PsiNamedElement getNamedElement(@NotNull final String the_name) {
232     // performance: check simple case first
233     PyExpression[] targets = getTargets();
234     if (targets.length == 1 && targets[0] instanceof PyTargetExpression) {
235       PyTargetExpression target = (PyTargetExpression)targets[0];
236       return !target.isQualified() && the_name.equals(target.getName()) ? target : null;
237     }
238     return PyUtil.IterHelper.findName(getNamedElements(), the_name);
239   }
240
241   @Override
242   public void subtreeChanged() {
243     super.subtreeChanged();
244     myTargets = null;
245   }
246 }