From f6fe43223b11cf6f5889b7573670032be607cabe Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Thu, 7 Sep 2017 18:18:06 +0300
Subject: [PATCH] PY-20401 Rehighlight None, True and False to keep their color
 inside annotations

Otherwise, they have the same color as enclosing annotation that is
different from the way we treat builtin symbols like 'str' and 'int'
there.
---
 .../src/com/jetbrains/python/PyTokenTypes.java       |  1 +
 .../python/validation/HighlightingAnnotator.java     | 12 ++++++++++++
 python/testData/highlighting/annotations.py          |  2 +-
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/python/psi-api/src/com/jetbrains/python/PyTokenTypes.java b/python/psi-api/src/com/jetbrains/python/PyTokenTypes.java
index 0e93006168e1..63b80160a75b 100644
--- a/python/psi-api/src/com/jetbrains/python/PyTokenTypes.java
+++ b/python/psi-api/src/com/jetbrains/python/PyTokenTypes.java
@@ -169,6 +169,7 @@ public class PyTokenTypes {
   public static final TokenSet NUMERIC_LITERALS = TokenSet.create(FLOAT_LITERAL, INTEGER_LITERAL, IMAGINARY_LITERAL);
   public static final TokenSet BOOL_LITERALS = TokenSet.create(TRUE_KEYWORD, FALSE_KEYWORD);
   public static final TokenSet SCALAR_LITERALS = TokenSet.orSet(BOOL_LITERALS, NUMERIC_LITERALS, TokenSet.create(NONE_KEYWORD));
+  public static final TokenSet EXPRESSION_KEYWORDS = TokenSet.create(TRUE_KEYWORD, FALSE_KEYWORD, NONE_KEYWORD);
 
   public static final TokenSet AUG_ASSIGN_OPERATIONS = TokenSet.create(PLUSEQ, MINUSEQ, MULTEQ, ATEQ, DIVEQ,
       PERCEQ, EXPEQ, GTGTEQ, LTLTEQ, ANDEQ, OREQ, XOREQ, FLOORDIVEQ);
diff --git a/python/src/com/jetbrains/python/validation/HighlightingAnnotator.java b/python/src/com/jetbrains/python/validation/HighlightingAnnotator.java
index 488d27ac6859..1dadf823d2ca 100644
--- a/python/src/com/jetbrains/python/validation/HighlightingAnnotator.java
+++ b/python/src/com/jetbrains/python/validation/HighlightingAnnotator.java
@@ -17,7 +17,9 @@ package com.jetbrains.python.validation;
 
 import com.intellij.lang.ASTNode;
 import com.intellij.openapi.editor.colors.TextAttributesKey;
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.highlighting.PyHighlighter;
 import com.jetbrains.python.psi.*;
 
@@ -78,4 +80,14 @@ public class HighlightingAnnotator extends PyAnnotator {
       addHighlightingAnnotation(value, PyHighlighter.PY_ANNOTATION);
     }
   }
+
+  @Override
+  public void visitElement(PsiElement element) {
+    // Highlight None, True and False as keywords once again inside annotations after PyHighlighter
+    // to keep their original color
+    if (PyTokenTypes.EXPRESSION_KEYWORDS.contains(element.getNode().getElementType()) &&
+        PsiTreeUtil.getParentOfType(element, PyAnnotation.class) != null) {
+      addHighlightingAnnotation(element, PyHighlighter.PY_KEYWORD);
+    }
+  }
 }
diff --git a/python/testData/highlighting/annotations.py b/python/testData/highlighting/annotations.py
index 9fa0e365b05d..b2fa48600ae5 100644
--- a/python/testData/highlighting/annotations.py
+++ b/python/testData/highlighting/annotations.py
@@ -2,5 +2,5 @@ TOP_LEVEL: <info descr="PY.ANNOTATION"><info descr="PY.BUILTIN_NAME">str</info><
 
 class <info descr="PY.CLASS_DEFINITION">C</info>:
     attr: <info descr="PY.ANNOTATION">Optional[Any]</info> = None
-    def <info descr="PY.FUNC_DEFINITION">method</info>(<info descr="PY.SELF_PARAMETER">self</info>, <info descr="PY.PARAMETER">xs</info>: <info descr="PY.ANNOTATION">List[<info descr="PY.BUILTIN_NAME">int</info>]</info>) -> <info descr="PY.ANNOTATION">None</info>:
+    def <info descr="PY.FUNC_DEFINITION">method</info>(<info descr="PY.SELF_PARAMETER">self</info>, <info descr="PY.PARAMETER">xs</info>: <info descr="PY.ANNOTATION">List[<info descr="PY.BUILTIN_NAME">int</info>]</info>) -> <info descr="PY.ANNOTATION"><info descr="PY.KEYWORD">None</info></info>:
         pass