python/src/com/jetbrains/python/validation/PyDefinitionsAnnotator.java

   1 /*
   2  * Copyright 2000-2017 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.validation;
  17
  18 import com.intellij.lang.ASTNode;
  19 import com.intellij.openapi.project.IndexNotReadyException;
  20 import com.intellij.psi.PsiElement;
  21 import com.intellij.psi.util.PsiTreeUtil;
  22 import com.jetbrains.python.PyNames;
  23 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
  24 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
  25 import com.jetbrains.python.highlighting.PyHighlighter;
  26 import com.jetbrains.python.psi.*;
  27 import org.jetbrains.annotations.NotNull;
  28
  29 import java.util.Objects;
  30
  31 /**
  32  * Highlights class definitions, functrion definitions, and decorators.
  33  * User: dcheryasov
  34  */
  35 public class PyDefinitionsAnnotator extends PyAnnotator {
  36
  37   @Override
  38   public void visitPyClass(PyClass node) {
  39     final ASTNode name_node = node.getNameNode();
  40     if (name_node != null) {
  41       addHighlightingAnnotation(name_node, PyHighlighter.PY_CLASS_DEFINITION);
  42     }
  43   }
  44
  45   @Override
  46   public void visitPyFunction(PyFunction node) {
  47     ASTNode nameNode = node.getNameNode();
  48     if (nameNode != null) {
  49       final String name = node.getName();
  50       LanguageLevel languageLevel = LanguageLevel.forElement(node);
  51       if (PyNames.UNDERSCORED_ATTRIBUTES.contains(name) || PyNames.getBuiltinMethods(languageLevel).containsKey(name)) {
  52         PyClass cls = node.getContainingClass();
  53         if (PyUtil.isNewMethod(node)) {
  54           boolean new_style_class = false;
  55           try {
  56             if (cls != null) new_style_class = cls.isNewStyleClass(null);
  57           }
  58           catch (IndexNotReadyException ignored) {
  59           }
  60           if (new_style_class) {
  61             addHighlightingAnnotation(nameNode, PyHighlighter.PY_PREDEFINED_DEFINITION);
  62           }
  63         }
  64         else {
  65           addHighlightingAnnotation(nameNode, PyHighlighter.PY_PREDEFINED_DEFINITION);
  66         }
  67       }
  68       else {
  69         if (ScopeUtil.getScopeOwner(node) instanceof PyFunction) {
  70           addHighlightingAnnotation(nameNode, PyHighlighter.PY_NESTED_FUNC_DEFINITION);
  71         }
  72         else {
  73           addHighlightingAnnotation(nameNode, PyHighlighter.PY_FUNC_DEFINITION);
  74         }
  75       }
  76     }
  77   }
  78
  79   @Override
  80   public void visitPyDecoratorList(PyDecoratorList node) {
  81     PyDecorator[] decos = node.getDecorators();
  82     for (PyDecorator deco : decos) {
  83       highlightDecorator(deco);
  84     }
  85   }
  86
  87   private void highlightDecorator(@NotNull PyDecorator node) {
  88     final PsiElement atSign = node.getFirstChild();
  89     if (atSign != null) {
  90       addHighlightingAnnotation(atSign, PyHighlighter.PY_DECORATOR);
  91       if (node.getQualifiedName() != null) {
  92         addHighlightingAnnotation(Objects.requireNonNull(node.getCallee()), PyHighlighter.PY_DECORATOR);
  93       }
  94     }
  95   }
  96 }