From a2ab96d619694f25b8af5999a19cd906e26daccf Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Wed, 6 Feb 2013 21:08:57 +0100
Subject: [PATCH] PyParameter.isSelf(); highlight first argument of
 @classmethod as self (PY-6881)

---
 .../com/jetbrains/python/psi/PyParameter.java |  6 ++++
 .../PyUnusedLocalInspectionVisitor.java       | 12 ++-----
 .../python/psi/impl/PyNamedParameterImpl.java | 29 +++++++++++++++--
 .../psi/impl/PySingleStarParameterImpl.java   |  5 +++
 .../python/psi/impl/PyTupleParameterImpl.java |  5 +++
 .../validation/HighlightingAnnotator.java     | 31 ++-----------------
 6 files changed, 48 insertions(+), 40 deletions(-)

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyParameter.java b/python/psi-api/src/com/jetbrains/python/psi/PyParameter.java
index 123505596fca..a63d47be38f5 100644
--- a/python/psi-api/src/com/jetbrains/python/psi/PyParameter.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyParameter.java
@@ -26,4 +26,10 @@ public interface PyParameter extends PyElement {
   PyExpression getDefaultValue();
 
   boolean hasDefaultValue();
+
+  /**
+   * @return true if the parameter is the 'self' parameter of an instance attribute function or a function
+   * annotated with @classmethod
+   */
+  boolean isSelf();
 }
diff --git a/python/src/com/jetbrains/python/inspections/PyUnusedLocalInspectionVisitor.java b/python/src/com/jetbrains/python/inspections/PyUnusedLocalInspectionVisitor.java
index 1454a9c19f42..5aae8e6ec16b 100644
--- a/python/src/com/jetbrains/python/inspections/PyUnusedLocalInspectionVisitor.java
+++ b/python/src/com/jetbrains/python/inspections/PyUnusedLocalInspectionVisitor.java
@@ -247,15 +247,9 @@ public class PyUnusedLocalInspectionVisitor extends PyInspectionVisitor {
                                             ? (PyNamedParameter) element
                                             : (PyNamedParameter) element.getParent();
           name = namedParameter.getName();
-          if (PsiTreeUtil.getParentOfType(element, PyClass.class) != null) {
-            // When function is inside a class, first parameter may be either self or cls which is always 'used'.
-            final PyFunction method = PsiTreeUtil.getParentOfType(element, PyFunction.class);
-            if (method != null && method.getModifier() != PyFunction.Modifier.STATICMETHOD) {
-              final PsiElement parent = namedParameter.getParent();
-              if (parent instanceof PyParameterList && ((PyParameterList)parent).getParameters()[0] == namedParameter) {
-                continue;
-              }
-            }
+          // When function is inside a class, first parameter may be either self or cls which is always 'used'.
+          if (namedParameter.isSelf()) {
+            continue;
           }
           if (myIgnoreLambdaParameters && PsiTreeUtil.getParentOfType(element, Callable.class) instanceof PyLambdaExpression) {
             continue;
diff --git a/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java b/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
index 7f16c71e37a3..1ebeec40f793 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
@@ -9,6 +9,7 @@ import com.intellij.psi.stubs.IStubElementType;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.IncorrectOperationException;
 import com.intellij.util.PlatformIcons;
+import com.jetbrains.cython.psi.CythonClass;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyTokenTypes;
@@ -157,11 +158,10 @@ public class PyNamedParameterImpl extends PyPresentableElementImpl<PyNamedParame
     final PsiElement parent = getStubOrPsiParent();
     if (parent instanceof PyParameterList) {
       PyParameterList parameterList = (PyParameterList)parent;
-      final PyParameter[] params = parameterList.getParameters();
       PyFunction func = parameterList.getContainingFunction();
       if (func != null) {
         final PyFunction.Modifier modifier = func.getModifier();
-        if (params [0] == this && modifier != PyFunction.Modifier.STATICMETHOD) {
+        if (isSelf()) {
           // must be 'self' or 'cls'
           final PyClass containingClass = func.getContainingClass();
           if (containingClass != null) {
@@ -239,4 +239,29 @@ public class PyNamedParameterImpl extends PyPresentableElementImpl<PyNamedParame
     }
     return new LocalSearchScope(getContainingFile());
   }
+
+  @Override
+  public boolean isSelf() {
+    if (isPositionalContainer() || isKeywordContainer()) {
+      return false;
+    }
+    PyFunction function = getStubOrPsiParentOfType(PyFunction.class);
+    if (function == null) {
+      return false;
+    }
+    final PyClass cls = function.getContainingClass();
+    if (cls != null && function.getParameterList().getParameters()[0] == this) {
+      if (cls instanceof CythonClass && ((CythonClass)cls).isCppClass()) {
+        return false;
+      }
+      if (PyNames.NEW.equals(function.getName())) {
+        return true;
+      }
+      final PyFunction.Modifier modifier = function.getModifier();
+      if (modifier != PyFunction.Modifier.STATICMETHOD) {
+        return true;
+      }
+    }
+    return false;
+  }
 }
diff --git a/python/src/com/jetbrains/python/psi/impl/PySingleStarParameterImpl.java b/python/src/com/jetbrains/python/psi/impl/PySingleStarParameterImpl.java
index 409f98d2892e..77686ba28652 100644
--- a/python/src/com/jetbrains/python/psi/impl/PySingleStarParameterImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PySingleStarParameterImpl.java
@@ -43,4 +43,9 @@ public class PySingleStarParameterImpl extends PyPresentableElementImpl<PySingle
   public boolean hasDefaultValue() {
     return false;
   }
+
+  @Override
+  public boolean isSelf() {
+    return false;
+  }
 }
diff --git a/python/src/com/jetbrains/python/psi/impl/PyTupleParameterImpl.java b/python/src/com/jetbrains/python/psi/impl/PyTupleParameterImpl.java
index 34daf762d8ff..69831de40578 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyTupleParameterImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyTupleParameterImpl.java
@@ -60,4 +60,9 @@ public class PyTupleParameterImpl extends PyPresentableElementImpl<PyTupleParame
   public PyParameter[] getContents() {
     return getStubOrPsiChildren(PythonDialectsTokenSetProvider.INSTANCE.getParameterTokens(), new PyParameter[0]);
   }
+
+  @Override
+  public boolean isSelf() {
+    return false;
+  }
 }
diff --git a/python/src/com/jetbrains/python/validation/HighlightingAnnotator.java b/python/src/com/jetbrains/python/validation/HighlightingAnnotator.java
index 3e665a2e6be8..ce7bc9e85ebd 100644
--- a/python/src/com/jetbrains/python/validation/HighlightingAnnotator.java
+++ b/python/src/com/jetbrains/python/validation/HighlightingAnnotator.java
@@ -3,9 +3,6 @@ package com.jetbrains.python.validation;
 import com.intellij.lang.ASTNode;
 import com.intellij.lang.annotation.Annotation;
 import com.intellij.psi.util.PsiTreeUtil;
-import com.intellij.util.ArrayUtil;
-import com.jetbrains.cython.psi.CythonClass;
-import com.jetbrains.python.PyNames;
 import com.jetbrains.python.highlighting.PyHighlighter;
 import com.jetbrains.python.psi.*;
 
@@ -18,7 +15,7 @@ public class HighlightingAnnotator extends PyAnnotator {
     PyFunction function = PsiTreeUtil.getParentOfType(node, PyFunction.class);
     if (function != null) {
       Annotation annotation = getHolder().createInfoAnnotation(node, null);
-      annotation.setTextAttributes(isSelf(node, function) ? PyHighlighter.PY_SELF_PARAMETER : PyHighlighter.PY_PARAMETER);
+      annotation.setTextAttributes(node.isSelf() ? PyHighlighter.PY_SELF_PARAMETER : PyHighlighter.PY_PARAMETER);
     }
   }
 
@@ -31,7 +28,7 @@ public class HighlightingAnnotator extends PyAnnotator {
         final PyNamedParameter element = function.getParameterList().findParameterByName(referencedName);
         if (element != null) {
           Annotation annotation = getHolder().createInfoAnnotation(node, null);
-          annotation.setTextAttributes(isSelf(element, function) ? PyHighlighter.PY_SELF_PARAMETER : PyHighlighter.PY_PARAMETER);
+          annotation.setTextAttributes(element.isSelf() ? PyHighlighter.PY_SELF_PARAMETER : PyHighlighter.PY_PARAMETER);
         }
       }
     }
@@ -45,28 +42,4 @@ public class HighlightingAnnotator extends PyAnnotator {
       annotation.setTextAttributes(PyHighlighter.PY_KEYWORD_ARGUMENT);
     }
   }
-
-  private static boolean isSelf(PyParameter node, PyFunction function) {
-    final PyNamedParameter named = node.getAsNamed();
-    if (named != null) {
-      if (named.isPositionalContainer() || named.isKeywordContainer()) {
-        return false;
-      }
-    }
-    final int index = ArrayUtil.find(function.getParameterList().getParameters(), node);
-    final PyClass cls = function.getContainingClass();
-    if (cls != null && index == 0) {
-      if (cls instanceof CythonClass && ((CythonClass)cls).isCppClass()) {
-        return false;
-      }
-      if (PyNames.NEW.equals(function.getName())) {
-        return true;
-      }
-      final PyFunction.Modifier modifier = function.getModifier();
-      if (modifier != PyFunction.Modifier.CLASSMETHOD && modifier != PyFunction.Modifier.STATICMETHOD) {
-        return true;
-      }
-    }
-    return false;
-  }
 }