PY-32302 Provide custom syntax highlighting for Python local variables

GitOrigin-RevId: af7099a6ba948c740c8183674b24eeeddf164e70

python/pluginResources/messages/PyBundle.properties

@@ -871,6 +871,7 @@ python.colors.string.binary.bytes=String//Binary (bytes)
 python.colors.line.comment=Line Comment
 python.colors.keyword=Keyword
 python.colors.number=Number
+python.colors.local.variables=Local variables
 
 python.new.project.synchronization.not.configured.dialog.title=Synchronization not Configured
 python.new.project.synchronization.not.configured.dialog.message=Local/Remote synchronization is not configured correctly.\n{0}\n\

python/python-psi-impl/src/com/jetbrains/python/highlighting/PyHighlighter.java

@@ -119,6 +119,7 @@ public class PyHighlighter extends SyntaxHighlighterBase {
   public static final TextAttributesKey PY_FSTRING_FRAGMENT_BRACES = TextAttributesKey.createTextAttributesKey("PY.FSTRING_FRAGMENT_BRACES", VALID_STRING_ESCAPE);
   public static final TextAttributesKey PY_FSTRING_FRAGMENT_COLON = TextAttributesKey.createTextAttributesKey("PY.FSTRING_FRAGMENT_COLON", VALID_STRING_ESCAPE);
   public static final TextAttributesKey PY_FSTRING_FRAGMENT_TYPE_CONVERSION = TextAttributesKey.createTextAttributesKey("PY.FSTRING_FRAGMENT_TYPE_CONVERSION", VALID_STRING_ESCAPE);
+  public static final TextAttributesKey PY_LOCAL_VARIABLE = TextAttributesKey.createTextAttributesKey("PY.LOCAL_VARIABLE", LOCAL_VARIABLE);
 
   /**
    * The 'heavy' constructor that initializes everything. PySyntaxHighlighterFactory caches such instances per level.

python/python-psi-impl/src/com/jetbrains/python/validation/PyVariableAnnotator.kt

@@ -0,0 +1,54 @@
+package com.jetbrains.python.validation
+
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil
+import com.jetbrains.python.highlighting.PyHighlighter
+import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.resolve.PyResolveUtil
+
+class PyVariableAnnotator : PyAnnotator() {
+
+  override fun visitPyTargetExpression(node: PyTargetExpression) {
+    if (node.isNonLocalOrGlobal() || node.isQualified) return
+
+    val scopeOwner = ScopeUtil.getScopeOwner(node)
+    if (node.parent is PyAssignmentStatement && scopeOwner is PyFunction) {
+      node.nameElement?.let { addHighlightingAnnotation(it, PyHighlighter.PY_LOCAL_VARIABLE) }
+    }
+  }
+
+  override fun visitPyReferenceExpression(node: PyReferenceExpression) {
+    PyResolveUtil.resolveLocally(node)
+      .filterIsInstance(PyTargetExpression::class.java)
+      .forEach { expression ->
+        if (ScopeUtil.getScopeOwner(expression) is PyFunction && !expression.isNonLocalOrGlobal()) {
+          addHighlightingAnnotation(node.node, PyHighlighter.PY_LOCAL_VARIABLE)
+        }
+      }
+  }
+
+  private fun PyQualifiedExpression.isNonLocalOrGlobal(): Boolean {
+    val qName = this.asQualifiedName()
+    val scopeOwner = ScopeUtil.getScopeOwner(this)
+
+    if (qName != null && scopeOwner is PyFunction) {
+      val scopesToLookUp = mutableListOf(scopeOwner)
+
+      scopesToLookUp.addAll(PsiTreeUtil.findChildrenOfType(scopeOwner, PyFunction::class.java))
+      scopesToLookUp.forEach { scope ->
+        if (PyResolveUtil.resolveLocally(scope, qName.toString()).containsNonLocalOrGlobal()) {
+          return true
+        }
+      }
+    }
+    return false
+  }
+
+  private fun Collection<PsiElement>.containsNonLocalOrGlobal(): Boolean =
+    this.filterIsInstance(PyTargetExpression::class.java)
+      .any { expression -> expression.parent.isNonLocalOrGlobal() }
+
+  private fun PsiElement.isNonLocalOrGlobal(): Boolean =
+    this is PyNonlocalStatement || this is PyGlobalStatement
+}

python/src/META-INF/python-core-common.xml

@@ -632,6 +632,7 @@
     <pyAnnotator implementation="com.jetbrains.python.validation.PyHighlightingAnnotator"/>
     <pyAnnotator implementation="com.jetbrains.python.validation.PyPatternAnnotator"/>
     <pyAnnotator implementation="com.jetbrains.python.validation.PyTryExceptAnnotator"/>
+    <pyAnnotator implementation="com.jetbrains.python.validation.PyVariableAnnotator"/>
 
     <!--stdlib-->
     <documentationLinkProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibDocumentationLinkProvider"/>

python/src/com/jetbrains/python/highlighting/PythonColorsPage.java

@@ -5,7 +5,6 @@ import com.google.common.collect.ImmutableMap;
 import com.intellij.application.options.colors.InspectionColorSettingsPage;
 import com.intellij.codeHighlighting.RainbowHighlighter;
 import com.intellij.lang.Language;
-import com.intellij.openapi.editor.DefaultLanguageHighlighterColors;
 import com.intellij.openapi.editor.colors.TextAttributesKey;
 import com.intellij.openapi.fileTypes.SyntaxHighlighter;
 import com.intellij.openapi.fileTypes.SyntaxHighlighterFactory;
@@ -68,6 +67,7 @@ public class PythonColorsPage implements RainbowColorSettingsPage, InspectionCol
     new AttributesDescriptor(PyBundle.message("python.colors.decorator"), PyHighlighter.PY_DECORATOR),
     new AttributesDescriptor(PyBundle.message("python.colors.class.definition"), PyHighlighter.PY_CLASS_DEFINITION),
     new AttributesDescriptor(PyBundle.message("python.colors.type.annotation"), PyHighlighter.PY_ANNOTATION),
+    new AttributesDescriptor(PyBundle.message("python.colors.local.variables"), PyHighlighter.PY_LOCAL_VARIABLE),
   };
 
   @NonNls private static final Map<String,TextAttributesKey> ourTagToDescriptorMap = ImmutableMap.<String, TextAttributesKey>builder()
@@ -86,7 +86,7 @@ public class PythonColorsPage implements RainbowColorSettingsPage, InspectionCol
     .put("call", PyHighlighter.PY_FUNCTION_CALL)
     .put("mcall", PyHighlighter.PY_METHOD_CALL)
     .put("annotation", PyHighlighter.PY_ANNOTATION)
-    .put("localVar", DefaultLanguageHighlighterColors.LOCAL_VARIABLE)
+    .put("localVar", PyHighlighter.PY_LOCAL_VARIABLE)
     .putAll(RainbowHighlighter.createRainbowHLM())
     .build();
 

python/testData/highlighting/localVariables.py

@@ -0,0 +1,8 @@
+def <info descr="PY.FUNC_DEFINITION">fun</info>():
+    <info descr="PY.LOCAL_VARIABLE">local_var</info> = "hello"
+    <info descr="PY.BUILTIN_NAME">print</info>(<info descr="PY.LOCAL_VARIABLE">local_var</info>)
+
+    def <info descr="PY.NESTED_FUNC_DEFINITION">nested</info>():
+        <info descr="PY.BUILTIN_NAME">print</info>(<info descr="PY.LOCAL_VARIABLE">local_var</info>)
+
+    return

python/testData/highlighting/nestedParamHighlightingInInnerFunc.py

@@ -1,8 +1,8 @@
 def <info descr="PY.FUNC_DEFINITION">outer_func</info>(<info descr="PY.PARAMETER">a</info>, <info descr="PY.PARAMETER">b</info>):
     def <info descr="PY.NESTED_FUNC_DEFINITION">inner_func_one</info>(<info descr="PY.PARAMETER">c</info>):
         def <info descr="PY.NESTED_FUNC_DEFINITION">inner_func_two</info>(<info descr="PY.PARAMETER">d</info>):
-            x = 10
-            return <info descr="PY.PARAMETER">a</info> + <info descr="PY.PARAMETER">b</info> + <info descr="PY.PARAMETER">c</info> + <info descr="PY.PARAMETER">d</info> + x
+            <info descr="PY.LOCAL_VARIABLE">x</info> = 10
+            return <info descr="PY.PARAMETER">a</info> + <info descr="PY.PARAMETER">b</info> + <info descr="PY.PARAMETER">c</info> + <info descr="PY.PARAMETER">d</info> + <info descr="PY.LOCAL_VARIABLE">x</info>
 
         return <info descr="PY.FUNCTION_CALL">inner_func_two</info>(4)
 

python/testData/highlighting/variableAnnotatedWithGlobalNotHighlightedAsLocal.py

@@ -0,0 +1,5 @@
+def <info descr="PY.FUNC_DEFINITION">foo</info>():
+  global g
+  g = "world!"
+
+<info descr="PY.BUILTIN_NAME">print</info>("Hello, " + g)

python/testData/highlighting/variableAnnotatedWithNonLocalNotHighlightedAsLocal.py

@@ -0,0 +1,7 @@
+def <info descr="PY.FUNC_DEFINITION">outer</info>():
+  x = "John"
+  def <info descr="PY.NESTED_FUNC_DEFINITION">inner</info>():
+    nonlocal x
+    x = "hello"
+  <info descr="PY.FUNCTION_CALL">inner</info>()
+  return x

python/testSrc/com/jetbrains/python/PythonHighlightingTest.java

@@ -565,6 +565,21 @@ public class PythonHighlightingTest extends PyTestCase {
     doTest(LanguageLevel.getLatest(), false, true);
   }
 
+  // PY-32302
+  public void testLocalVariables() {
+    doTest(LanguageLevel.getLatest(), false, true);
+  }
+
+  // PY-32302
+  public void testVariableAnnotatedWithNonLocalNotHighlightedAsLocal() {
+    doTest(LanguageLevel.getLatest(), false, true);
+  }
+
+  // PY-32302
+  public void testVariableAnnotatedWithGlobalNotHighlightedAsLocal() {
+    doTest(LanguageLevel.getLatest(), false, true);
+  }
+
   @NotNull
   private static EditorColorsScheme createTemporaryColorScheme() {
     EditorColorsManager manager = EditorColorsManager.getInstance();