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 }