PY-24653 Fix `Self` highlighting in nested functions

(cherry picked from commit fbfc5683399723ac0d49e5aa337024154dd6d3a9)

IJ-CR-16991

GitOrigin-RevId: 69084562247b0105f268fac9f442e02756240173

diff --git a/python/src/com/jetbrains/python/validation/HighlightingAnnotator.java b/python/src/com/jetbrains/python/validation/HighlightingAnnotator.java
index b48a0a32c61cab0f941065526c0f5f749b7f47dd..2416db4fb292ba5216915ebbc72ee479c9c3421e 100644 (file)
--- a/python/src/com/jetbrains/python/validation/HighlightingAnnotator.java
+++ b/python/src/com/jetbrains/python/validation/HighlightingAnnotator.java
@@ -6,11 +6,14 @@ import com.intellij.lang.annotation.HighlightSeverity;
 import com.intellij.openapi.editor.colors.TextAttributesKey;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.ObjectUtils;
 import com.jetbrains.python.PyTokenTypes;
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.highlighting.PyHighlighter;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import static com.jetbrains.python.psi.PyUtil.as;
 
@@ -46,13 +49,19 @@ public class HighlightingAnnotator extends PyAnnotator {
   public void visitPyReferenceExpression(@NotNull PyReferenceExpression node) {
     final String referencedName = node.getReferencedName();
     if (!node.isQualified() && referencedName != null) {
-      final PyFunction function = PsiTreeUtil.getParentOfType(node, PyFunction.class);
+      PyFunction function = PsiTreeUtil.getParentOfType(node, PyFunction.class);
       if (function != null) {
-        final PyNamedParameter element = function.getParameterList().findParameterByName(referencedName);
+        PyNamedParameter element = function.getParameterList().findParameterByName(referencedName);
         if (element != null) {
           final TextAttributesKey attrKey = element.isSelf() ? PyHighlighter.PY_SELF_PARAMETER : PyHighlighter.PY_PARAMETER;
           addHighlightingAnnotation(node, attrKey);
         }
+        else {
+          element = findParameterRecursively(function, referencedName);
+          if (element != null && element.isSelf()) {
+            addHighlightingAnnotation(node, PyHighlighter.PY_SELF_PARAMETER);
+          }
+        }
       }
     }
   }
@@ -98,4 +107,15 @@ public class HighlightingAnnotator extends PyAnnotator {
       addHighlightingAnnotation(element, PyHighlighter.PY_KEYWORD);
     }
   }
+
+  @Nullable
+  private static PyNamedParameter findParameterRecursively(@NotNull PyFunction function, @NotNull String referencedName) {
+    for (PyFunction f = function; f != null; f = ObjectUtils.tryCast(ScopeUtil.getScopeOwner(f), PyFunction.class)) {
+      PyNamedParameter element = f.getParameterList().findParameterByName(referencedName);
+      if (element != null) {
+        return element;
+      }
+    }
+    return null;
+  }
 }

diff --git a/python/testData/highlighting/selfHighlightingInInnerFunc.py b/python/testData/highlighting/selfHighlightingInInnerFunc.py
new file mode 100644 (file)
index 0000000..43b7844
--- /dev/null
+++ b/python/testData/highlighting/selfHighlightingInInnerFunc.py
@@ -0,0 +1,21 @@
+class <info descr="PY.CLASS_DEFINITION">ExampleClass</info>(<info descr="PY.BUILTIN_NAME">object</info>):
+    def <info descr="PY.PREDEFINED_DEFINITION">__init__</info>(<info descr="PY.SELF_PARAMETER">self</info>):
+        <info descr="PY.SELF_PARAMETER">self</info>.ex1 = 1
+
+    def <info descr="PY.FUNC_DEFINITION">outer_test_func</info>(<info descr="PY.SELF_PARAMETER">self</info>):
+        def <info descr="PY.NESTED_FUNC_DEFINITION">inner_test_func</info>():
+            def <info descr="PY.NESTED_FUNC_DEFINITION">inner_test_func_lvl2</info>():
+                <info descr="PY.SELF_PARAMETER">self</info>.ex1 = 10
+
+            <info descr="PY.SELF_PARAMETER">self</info>.ex1 = 2
+
+        return inner_test_func
+
+    def <info descr="PY.FUNC_DEFINITION">outer_test_func_with_non_default_self_name</info>(<info descr="PY.SELF_PARAMETER">this</info>):
+            def <info descr="PY.NESTED_FUNC_DEFINITION">inner_test_func</info>():
+                def <info descr="PY.NESTED_FUNC_DEFINITION">inner_test_func_lvl2</info>():
+                    <info descr="PY.SELF_PARAMETER">this</info>.ex1 = 10
+
+                <info descr="PY.SELF_PARAMETER">this</info>.ex1 = 2
+
+            return inner_test_func
\ No newline at end of file

diff --git a/python/testSrc/com/jetbrains/python/PythonHighlightingTest.java b/python/testSrc/com/jetbrains/python/PythonHighlightingTest.java
index b55129b2f81bab1dbd9cff63e202c006c9c1f833..3320af0ab9c782065c2e4b34c742dc031b465184 100644 (file)
--- a/python/testSrc/com/jetbrains/python/PythonHighlightingTest.java
+++ b/python/testSrc/com/jetbrains/python/PythonHighlightingTest.java
@@ -533,6 +533,11 @@ public class PythonHighlightingTest extends PyTestCase {
     doTest(LanguageLevel.PYTHON39, false, false);
   }
 
+  // PY-24653
+  public void testSelfHighlightingInInnerFunc() {
+    doTest(LanguageLevel.getLatest(), false, true);
+  }
+
   @NotNull
   private static EditorColorsScheme createTemporaryColorScheme() {
     EditorColorsManager manager = EditorColorsManager.getInstance();