496caeac594110556db5e78d8b0710db5cf0472e
[idea/community.git] / python / psi-api / src / com / jetbrains / python / psi / impl / PyPsiUtils.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.base.Preconditions;
19 import com.intellij.lang.ASTNode;
20 import com.intellij.lang.injection.InjectedLanguageManager;
21 import com.intellij.openapi.diagnostic.Logger;
22 import com.intellij.openapi.module.Module;
23 import com.intellij.openapi.util.Couple;
24 import com.intellij.openapi.util.text.StringUtil;
25 import com.intellij.psi.*;
26 import com.intellij.psi.stubs.StubElement;
27 import com.intellij.psi.tree.IElementType;
28 import com.intellij.psi.tree.TokenSet;
29 import com.intellij.psi.util.PsiTreeUtil;
30 import com.intellij.psi.util.PsiUtilCore;
31 import com.intellij.psi.util.QualifiedName;
32 import com.intellij.util.containers.ContainerUtil;
33 import com.jetbrains.python.PyTokenTypes;
34 import com.jetbrains.python.psi.*;
35 import org.jetbrains.annotations.NotNull;
36 import org.jetbrains.annotations.Nullable;
37
38 import java.lang.reflect.Array;
39 import java.util.ArrayList;
40 import java.util.LinkedList;
41 import java.util.List;
42
43 /**
44  * @author max
45  */
46 public class PyPsiUtils {
47
48   private static final Logger LOG = Logger.getInstance(PyPsiUtils.class.getName());
49
50   private PyPsiUtils() {
51   }
52
53   @NotNull
54   protected static <T extends PyElement> T[] nodesToPsi(ASTNode[] nodes, T[] array) {
55     T[] psiElements = (T[])Array.newInstance(array.getClass().getComponentType(), nodes.length);
56     for (int i = 0; i < nodes.length; i++) {
57       //noinspection unchecked
58       psiElements[i] = (T)nodes[i].getPsi();
59     }
60     return psiElements;
61   }
62
63   /**
64    * Finds the closest comma after the element skipping any whitespaces in-between.
65    */
66   @Nullable
67   public static PsiElement getPrevComma(@NotNull PsiElement element) {
68     final PsiElement prevNode = getPrevNonWhitespaceSibling(element);
69     return prevNode != null && prevNode.getNode().getElementType() == PyTokenTypes.COMMA ? prevNode : null;
70   }
71
72   /**
73    * Finds first non-whitespace sibling before given PSI element.
74    */
75   @Nullable
76   public static PsiElement getPrevNonWhitespaceSibling(@Nullable PsiElement element) {
77     return PsiTreeUtil.skipSiblingsBackward(element, PsiWhiteSpace.class);
78   }
79
80   /**
81    * Finds first non-whitespace sibling before given AST node.
82    */
83   @Nullable
84   public static ASTNode getPrevNonWhitespaceSibling(@NotNull ASTNode node) {
85     return skipSiblingsBackward(node, TokenSet.create(TokenType.WHITE_SPACE));
86   }
87
88   /**
89    * Finds first sibling that is neither comment, nor whitespace before given element.
90    * @param strict prohibit returning element itself
91    */
92   @Nullable
93   public static PsiElement getPrevNonCommentSibling(@Nullable PsiElement start, boolean strict) {
94     if (!strict && !(start instanceof PsiWhiteSpace || start instanceof PsiComment)) {
95       return start;
96     }
97     return PsiTreeUtil.skipSiblingsBackward(start, PsiWhiteSpace.class, PsiComment.class);
98   }
99
100   /**
101    * Finds the closest comma after the element skipping any whitespaces in-between.
102    */
103   @Nullable
104   public static PsiElement getNextComma(@NotNull PsiElement element) {
105     final PsiElement nextNode = getNextNonWhitespaceSibling(element);
106     return nextNode != null && nextNode.getNode().getElementType() == PyTokenTypes.COMMA ? nextNode : null;
107   }
108
109   /**
110    * Finds first non-whitespace sibling after given PSI element.
111    */
112   @Nullable
113   public static PsiElement getNextNonWhitespaceSibling(@Nullable PsiElement element) {
114     return PsiTreeUtil.skipSiblingsForward(element, PsiWhiteSpace.class);
115   }
116
117   /**
118    * Finds first non-whitespace sibling after given PSI element but stops at first whitespace containing line feed.  
119    */
120   @Nullable
121   public static PsiElement getNextNonWhitespaceSiblingOnSameLine(@NotNull PsiElement element) {
122     PsiElement cur = element.getNextSibling();
123     while (cur != null) {
124       if (!(cur instanceof PsiWhiteSpace)) {
125         return cur;
126       }
127       else if (cur.textContains('\n')) {
128         break;
129       }
130       cur = cur.getNextSibling();
131     }
132     return null;
133   }
134   
135   /**
136    * Finds first non-whitespace sibling after given AST node.
137    */
138   @Nullable
139   public static ASTNode getNextNonWhitespaceSibling(@NotNull ASTNode after) {
140     return skipSiblingsForward(after, TokenSet.create(TokenType.WHITE_SPACE));
141   }
142
143   /**
144    * Finds first sibling that is neither comment, nor whitespace after given element.
145    * @param strict prohibit returning element itself
146    */
147   @Nullable
148   public static PsiElement getNextNonCommentSibling(@Nullable PsiElement start, boolean strict) {
149     if (!strict && !(start instanceof PsiWhiteSpace || start instanceof PsiComment)) {
150       return start;
151     }
152     return PsiTreeUtil.skipSiblingsForward(start, PsiWhiteSpace.class, PsiComment.class);
153   }
154
155   /**
156    * Finds first token after given element that doesn't consist solely of spaces and is not empty (e.g. error marker).
157    * @param ignoreComments ignore commentaries as well
158    */
159   @Nullable
160   public static PsiElement getNextSignificantLeaf(@Nullable PsiElement element, boolean ignoreComments) {
161     while (element != null && StringUtil.isEmptyOrSpaces(element.getText()) || ignoreComments && element instanceof PsiComment) {
162       element = PsiTreeUtil.nextLeaf(element);
163     }
164     return element;
165   }
166
167   /**
168    * Finds first token before given element that doesn't consist solely of spaces and is not empty (e.g. error marker).
169    * @param ignoreComments ignore commentaries as well
170    */
171   @Nullable
172   public static PsiElement getPrevSignificantLeaf(@Nullable PsiElement element, boolean ignoreComments) {
173     while (element != null && StringUtil.isEmptyOrSpaces(element.getText()) || ignoreComments && element instanceof PsiComment) {
174       element = PsiTreeUtil.prevLeaf(element);
175     }
176     return element;
177   }
178
179   /**
180    * Finds the closest comma looking for the next comma first and then for the preceding one.
181    */
182   @Nullable
183   public static PsiElement getAdjacentComma(@NotNull PsiElement element) {
184     final PsiElement nextComma = getNextComma(element);
185     return nextComma != null ? nextComma : getPrevComma(element);
186   }
187
188   /**
189    * Works similarly to {@link PsiTreeUtil#skipSiblingsForward(PsiElement, Class[])}, but for AST nodes.
190    */
191   @Nullable
192   public static ASTNode skipSiblingsForward(@Nullable ASTNode node, @NotNull TokenSet types) {
193     if (node == null) {
194       return null;
195     }
196     for (ASTNode next = node.getTreeNext(); next != null; next = next.getTreeNext()) {
197       if (!types.contains(next.getElementType())) {
198         return next;
199       }
200     }
201     return null;
202   }
203
204   /**
205    * Works similarly to {@link PsiTreeUtil#skipSiblingsBackward(PsiElement, Class[])}, but for AST nodes.
206    */
207   @Nullable
208   public static ASTNode skipSiblingsBackward(@Nullable ASTNode node, @NotNull TokenSet types) {
209     if (node == null) {
210       return null;
211     }
212     for (ASTNode prev = node.getTreePrev(); prev != null; prev = prev.getTreePrev()) {
213       if (!types.contains(prev.getElementType())) {
214         return prev;
215       }
216     }
217     return null;
218   }
219
220     /**
221    * Returns first child psi element with specified element type or {@code null} if no such element exists.
222    * Semantically it's the same as {@code getChildByFilter(element, TokenSet.create(type), 0)}.
223    *
224    * @param element tree parent node
225    * @param type    element type expected
226    * @return child element described
227    */
228   @Nullable
229   public static PsiElement getFirstChildOfType(@NotNull final PsiElement element, @NotNull PyElementType type) {
230     final ASTNode child = element.getNode().findChildByType(type);
231     return child != null ? child.getPsi() : null;
232   }
233
234   /**
235    * Returns child element in the psi tree
236    *
237    * @param filter  Types of expected child
238    * @param number  number
239    * @param element tree parent node
240    * @return PsiElement - child psiElement
241    */
242   @Nullable
243   public static PsiElement getChildByFilter(@NotNull PsiElement element, @NotNull TokenSet filter, int number) {
244     final ASTNode node = element.getNode();
245     if (node != null) {
246       final ASTNode[] children = node.getChildren(filter);
247       return (0 <= number && number < children.length) ? children[number].getPsi() : null;
248     }
249     return null;
250   }
251
252   public static void addBeforeInParent(@NotNull final PsiElement anchor, @NotNull final PsiElement... newElements) {
253     final ASTNode anchorNode = anchor.getNode();
254     LOG.assertTrue(anchorNode != null);
255     for (PsiElement newElement : newElements) {
256       anchorNode.getTreeParent().addChild(newElement.getNode(), anchorNode);
257     }
258   }
259
260   public static void removeElements(@NotNull final PsiElement... elements) {
261     final ASTNode parentNode = elements[0].getParent().getNode();
262     LOG.assertTrue(parentNode != null);
263     for (PsiElement element : elements) {
264       //noinspection ConstantConditions
265       parentNode.removeChild(element.getNode());
266     }
267   }
268
269   @Nullable
270   public static PsiElement getStatement(@NotNull final PsiElement element) {
271     final PyElement compStatement = getStatementList(element);
272     if (compStatement == null) {
273       return null;
274     }
275     return getParentRightBefore(element, compStatement);
276   }
277
278   public static PyElement getStatementList(final PsiElement element) {
279     //noinspection ConstantConditions
280     return element instanceof PyFile || element instanceof PyStatementList
281            ? (PyElement)element
282            : PsiTreeUtil.getParentOfType(element, PyFile.class, PyStatementList.class);
283   }
284
285   /**
286    * Returns ancestor of the element that is also direct child of the given super parent.
287    *
288    * @param element     element to start search from
289    * @param superParent direct parent of the desired ancestor
290    * @return described element or {@code null} if it doesn't exist
291    */
292   @Nullable
293   public static PsiElement getParentRightBefore(@NotNull PsiElement element, @NotNull final PsiElement superParent) {
294     return PsiTreeUtil.findFirstParent(element, false, element1 -> element1.getParent() == superParent);
295   }
296
297   public static List<PsiElement> collectElements(final PsiElement statement1, final PsiElement statement2) {
298     // Process ASTNodes here to handle all the nodes
299     final ASTNode node1 = statement1.getNode();
300     final ASTNode node2 = statement2.getNode();
301     final ASTNode parentNode = node1.getTreeParent();
302
303     boolean insideRange = false;
304     final List<PsiElement> result = new ArrayList<>();
305     for (ASTNode node : parentNode.getChildren(null)) {
306       // start
307       if (node1 == node) {
308         insideRange = true;
309       }
310       if (insideRange) {
311         result.add(node.getPsi());
312       }
313       // stop
314       if (node == node2) {
315         break;
316       }
317     }
318     return result;
319   }
320
321   public static int getElementIndentation(final PsiElement element) {
322     final PsiElement compStatement = getStatementList(element);
323     final PsiElement statement = getParentRightBefore(element, compStatement);
324     if (statement == null) {
325       return 0;
326     }
327     PsiElement sibling = statement.getPrevSibling();
328     if (sibling == null) {
329       sibling = compStatement.getPrevSibling();
330     }
331     final String whitespace = sibling instanceof PsiWhiteSpace ? sibling.getText() : "";
332     final int i = whitespace.lastIndexOf("\n");
333     return i != -1 ? whitespace.length() - i - 1 : 0;
334   }
335
336   public static void removeRedundantPass(final PyStatementList statementList) {
337     final PyStatement[] statements = statementList.getStatements();
338     if (statements.length > 1) {
339       for (PyStatement statement : statements) {
340         if (statement instanceof PyPassStatement) {
341           statement.delete();
342         }
343       }
344     }
345   }
346
347   public static boolean isMethodContext(final PsiElement element) {
348     final PsiNamedElement parent = PsiTreeUtil.getParentOfType(element, PyFile.class, PyFunction.class, PyClass.class);
349     // In case if element is inside method which is inside class
350     if (parent instanceof PyFunction && PsiTreeUtil.getParentOfType(parent, PyFile.class, PyClass.class) instanceof PyClass) {
351       return true;
352     }
353     return false;
354   }
355
356
357   @NotNull
358   public static PsiElement getRealContext(@NotNull final PsiElement element) {
359     assertValid(element);
360     final PsiFile file = element.getContainingFile();
361     if (file instanceof PyExpressionCodeFragment) {
362       final PsiElement context = file.getContext();
363       if (LOG.isDebugEnabled()) {
364         LOG.debug("PyPsiUtil.getRealContext(" + element + ") is called. Returned " + context + ". Element inside code fragment");
365       }
366       return context != null ? context : element;
367     }
368     else {
369       if (LOG.isDebugEnabled()) {
370         LOG.debug("PyPsiUtil.getRealContext(" + element + ") is called. Returned " + element + ".");
371       }
372       return element;
373     }
374   }
375
376   /**
377    * Removes comma closest to the given child node along with any whitespaces around. First following comma is checked and only
378    * then, if it doesn't exists, preceding one.
379    *
380    * @param element parent node
381    * @param child   child node comma should be adjacent to
382    * @see #getAdjacentComma(PsiElement)
383    */
384   public static void deleteAdjacentCommaWithWhitespaces(@NotNull PsiElement element, @NotNull PsiElement child) {
385     final PsiElement commaNode = getAdjacentComma(child);
386     if (commaNode != null) {
387       final PsiElement nextNonWhitespace = getNextNonWhitespaceSibling(commaNode);
388       final PsiElement last = nextNonWhitespace == null ? element.getLastChild() : nextNonWhitespace.getPrevSibling();
389       final PsiElement prevNonWhitespace = getPrevNonWhitespaceSibling(commaNode);
390       final PsiElement first = prevNonWhitespace == null ? element.getFirstChild() : prevNonWhitespace.getNextSibling();
391       element.deleteChildRange(first, last);
392     }
393   }
394
395   /**
396    * Returns comments preceding given elements as pair of the first and the last such comments. Comments should not be
397    * separated by any empty line.
398    * @param element element comments should be adjacent to
399    * @return described range or {@code null} if there are no such comments
400    */
401   @Nullable
402   public static Couple<PsiComment> getPrecedingComments(@NotNull PsiElement element) {
403     PsiComment firstComment = null, lastComment = null;
404     overComments:
405     while (true) {
406       int newLinesCount = 0;
407       for (element = element.getPrevSibling(); element instanceof PsiWhiteSpace; element = element.getPrevSibling()) {
408         newLinesCount += StringUtil.getLineBreakCount(element.getText());
409         if (newLinesCount > 1) {
410           break overComments;
411         }
412       }
413       if (element instanceof PsiComment) {
414         if (lastComment == null) {
415           lastComment = (PsiComment)element;
416         }
417         firstComment = (PsiComment)element;
418       }
419       else {
420         break;
421       }
422     }
423     return lastComment == null ? null : Couple.of(firstComment, lastComment);
424   }
425
426   @NotNull
427   static <T, U extends PsiElement> List<T> collectStubChildren(U e,
428                                                                final StubElement<U> stub, final IElementType elementType,
429                                                                final Class<T> itemClass) {
430     final List<T> result = new ArrayList<>();
431     if (stub != null) {
432       final List<StubElement> children = stub.getChildrenStubs();
433       for (StubElement child : children) {
434         if (child.getStubType() == elementType) {
435           //noinspection unchecked
436           result.add((T)child.getPsi());
437         }
438       }
439     }
440     else {
441       e.acceptChildren(new TopLevelVisitor() {
442         @Override
443         protected void checkAddElement(PsiElement node) {
444           if (node.getNode().getElementType() == elementType) {
445             //noinspection unchecked
446             result.add((T)node);
447           }
448         }
449       });
450     }
451     return result;
452   }
453
454   static List<PsiElement> collectAllStubChildren(PsiElement e, StubElement stub) {
455     final List<PsiElement> result = new ArrayList<>();
456     if (stub != null) {
457       final List<StubElement> children = stub.getChildrenStubs();
458       for (StubElement child : children) {
459         //noinspection unchecked
460         result.add(child.getPsi());
461       }
462     }
463     else {
464       e.acceptChildren(new TopLevelVisitor() {
465         @Override
466         protected void checkAddElement(PsiElement node) {
467           result.add(node);
468         }
469       });
470     }
471     return result;
472   }
473
474   public static int findArgumentIndex(PyCallExpression call, PsiElement argument) {
475     final PyExpression[] args = call.getArguments();
476     for (int i = 0; i < args.length; i++) {
477       PyExpression expression = args[i];
478       if (expression instanceof PyKeywordArgument) {
479         expression = ((PyKeywordArgument)expression).getValueExpression();
480       }
481       expression = flattenParens(expression);
482       if (expression == argument) {
483         return i;
484       }
485     }
486     return -1;
487   }
488
489   @Nullable
490   public static PyTargetExpression getAttribute(@NotNull final PyFile file, @NotNull final String name) {
491     PyTargetExpression attr = file.findTopLevelAttribute(name);
492     if (attr == null) {
493       for (PyFromImportStatement element : file.getFromImports()) {
494         PyReferenceExpression expression = element.getImportSource();
495         if (expression == null) continue;
496         final PsiElement resolved = expression.getReference().resolve();
497         if (resolved instanceof PyFile && resolved != file) {
498           return ((PyFile)resolved).findTopLevelAttribute(name);
499         }
500       }
501     }
502     return attr;
503   }
504
505   public static List<PyExpression> getAttributeValuesFromFile(@NotNull PyFile file, @NotNull String name) {
506     List<PyExpression> result = ContainerUtil.newArrayList();
507     final PyTargetExpression attr = file.findTopLevelAttribute(name);
508     if (attr != null) {
509       sequenceToList(result, attr.findAssignedValue());
510     }
511     return result;
512   }
513
514   public static void sequenceToList(List<PyExpression> result, PyExpression value) {
515     value = flattenParens(value);
516     if (value instanceof PySequenceExpression) {
517       result.addAll(ContainerUtil.newArrayList(((PySequenceExpression)value).getElements()));
518     }
519     else {
520       result.add(value);
521     }
522   }
523
524   public static List<String> getStringValues(PyExpression[] elements) {
525     List<String> results = ContainerUtil.newArrayList();
526     for (PyExpression element : elements) {
527       if (element instanceof PyStringLiteralExpression) {
528         results.add(((PyStringLiteralExpression)element).getStringValue());
529       }
530     }
531     return results;
532   }
533
534   @Nullable
535   public static PyExpression flattenParens(@Nullable PyExpression expr) {
536     while (expr instanceof PyParenthesizedExpression) {
537       expr = ((PyParenthesizedExpression)expr).getContainedExpression();
538     }
539     return expr;
540   }
541
542   @Nullable
543   public static String strValue(@Nullable PyExpression expression) {
544     return expression instanceof PyStringLiteralExpression ? ((PyStringLiteralExpression)expression).getStringValue() : null;
545   }
546
547   public static boolean isBefore(@NotNull final PsiElement element, @NotNull final PsiElement element2) {
548     // TODO: From RubyPsiUtil, should be moved to PsiTreeUtil
549     return element.getTextOffset() <= element2.getTextOffset();
550   }
551
552   @Nullable
553   public static QualifiedName asQualifiedName(@Nullable PyExpression expr) {
554     return expr instanceof PyQualifiedExpression ? ((PyQualifiedExpression)expr).asQualifiedName() : null;
555   }
556
557   @Nullable
558   public static PyExpression getFirstQualifier(@NotNull PyQualifiedExpression expr) {
559     final List<PyExpression> expressions = unwindQualifiers(expr);
560     if (!expressions.isEmpty()) {
561       return expressions.get(0);
562     }
563     return null;
564   }
565
566   @NotNull
567   public static String toPath(@Nullable PyQualifiedExpression expr) {
568     if (expr != null) {
569       final QualifiedName qName = expr.asQualifiedName();
570       if (qName != null) {
571         return qName.toString();
572       }
573       final String name = expr.getName();
574       if (name != null) {
575         return name;
576       }
577     }
578     return "";
579   }
580
581   @Nullable
582   protected static QualifiedName asQualifiedName(@NotNull PyQualifiedExpression expr) {
583     return fromReferenceChain(unwindQualifiers(expr));
584   }
585
586   @NotNull
587   private static List<PyExpression> unwindQualifiers(@NotNull final PyQualifiedExpression expr) {
588     final List<PyExpression> path = new LinkedList<>();
589     PyQualifiedExpression e = expr;
590     while (e != null) {
591       path.add(0, e);
592       final PyExpression q = e.getQualifier();
593       e = q instanceof PyQualifiedExpression ? (PyQualifiedExpression)q : null;
594     }
595     return path;
596   }
597
598   @Nullable
599   private static QualifiedName fromReferenceChain(@NotNull List<PyExpression> components) {
600     final List<String> componentNames = new ArrayList<>(components.size());
601     for (PyExpression component : components) {
602       final String refName = (component instanceof PyQualifiedExpression) ? ((PyQualifiedExpression)component).getReferencedName() : null;
603       if (refName == null) {
604         return null;
605       }
606       componentNames.add(refName);
607     }
608     return QualifiedName.fromComponents(componentNames);
609   }
610
611   /**
612    * Wrapper for {@link PsiUtilCore#ensureValid(PsiElement)} that skips nulls
613    */
614   public static void assertValid(@Nullable final PsiElement element) {
615     if (element == null) {
616       return;
617     }
618     PsiUtilCore.ensureValid(element);
619   }
620
621   public static void assertValid(@NotNull final Module module) {
622     Preconditions.checkArgument(!module.isDisposed(), String.format("Module %s is disposed", module));
623   }
624
625   @NotNull
626   public static PsiFileSystemItem getFileSystemItem(@NotNull PsiElement element) {
627     if (element instanceof PsiFileSystemItem) {
628       return (PsiFileSystemItem)element;
629     }
630     return element.getContainingFile();
631   }
632
633   private static abstract class TopLevelVisitor extends PyRecursiveElementVisitor {
634     public void visitPyElement(final PyElement node) {
635       super.visitPyElement(node);
636       checkAddElement(node);
637     }
638
639     public void visitPyClass(final PyClass node) {
640       checkAddElement(node);  // do not recurse into functions
641     }
642
643     public void visitPyFunction(final PyFunction node) {
644       checkAddElement(node);  // do not recurse into classes
645     }
646
647     protected abstract void checkAddElement(PsiElement node);
648   }
649
650   /**
651    * Returns text of the given PSI element. Unlike obvious {@link PsiElement#getText()} this method unescapes text of the element if latter
652    * belongs to injected code fragment using {@link InjectedLanguageManager#getUnescapedText(PsiElement)}.
653    *
654    * @param element PSI element which text is needed
655    * @return text of the element with any host escaping removed
656    */
657   @NotNull
658   public static String getElementTextWithoutHostEscaping(@NotNull PsiElement element) {
659     final InjectedLanguageManager manager = InjectedLanguageManager.getInstance(element.getProject());
660     if (manager.isInjectedFragment(element.getContainingFile())) {
661       return manager.getUnescapedText(element);
662     }
663     else {
664       return element.getText();
665     }
666   }
667 }