PY-1767: auto-add a correct first parameter to a method + tests.

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

@@ -43,6 +43,7 @@
     <typedHandler implementation="com.jetbrains.python.codeInsight.KeywordTypedHandler" id="pyCommaAfterKwd"/>
     <typedHandler implementation="com.jetbrains.django.lang.template.editor.DjangoBracesInterpolationTypedHandler" id="pyDjangoBraceInterpolation"/>
     <typedHandler implementation="com.jetbrains.python.editor.ColonTypedHandler" id="pyColonUnindent"/>
+    <typedHandler implementation="com.jetbrains.python.codeInsight.PyMethodNameTypedHandler" id="pyMethodNameTypedHandler"/>
 
     <stubIndex implementation="com.jetbrains.python.psi.stubs.PyClassNameIndex"/>
     <stubIndex implementation="com.jetbrains.python.psi.stubs.PyClassNameIndexInsensitive"/>

python/src/com/jetbrains/python/codeInsight/PyMethodNameTypedHandler.java

@@ -0,0 +1,69 @@
+package com.jetbrains.python.codeInsight;
+
+import com.intellij.codeInsight.editorActions.TypedHandlerDelegate;
+import com.intellij.lang.ASTNode;
+import com.intellij.openapi.editor.Document;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.editor.EditorModificationUtil;
+import com.intellij.openapi.fileTypes.FileType;
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.PsiDocumentManager;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.jetbrains.python.PyNames;
+import com.jetbrains.python.PyTokenTypes;
+import com.jetbrains.python.PythonFileType;
+import com.jetbrains.python.psi.PyFunction;
+import com.jetbrains.python.psi.PyStringLiteralExpression;
+import com.jetbrains.python.psi.PyUtil;
+
+/**
+ * Adds appropriate first parameter to a freshly-typed method declaration.
+ * <br/>
+ * User: dcheryasov
+ * Date: 11/29/10 12:44 AM
+ */
+public class PyMethodNameTypedHandler extends TypedHandlerDelegate {
+  @Override
+  public Result beforeCharTyped(char character, Project project, Editor editor, PsiFile file, FileType fileType) {
+    if (!(fileType instanceof PythonFileType)) return Result.CONTINUE; // else we'd mess up with other file types!
+    if (character == '(') {
+      final Document document = editor.getDocument();
+      final PsiDocumentManager documentManager = PsiDocumentManager.getInstance(project);
+      final int offset = editor.getCaretModel().getOffset();
+
+      PsiElement token = file.findElementAt(offset - 1);
+      if (token == null) return Result.CONTINUE; // sanity check: beyond EOL
+
+      final ASTNode token_node = token.getNode();
+      if (token_node != null && token_node.getElementType() == PyTokenTypes.IDENTIFIER) {
+        PsiElement maybe_def = PyUtil.getFirstNonCommentBefore(token.getPrevSibling());
+        if (maybe_def != null) {
+          ASTNode def_node = maybe_def.getNode();
+          if (def_node != null && def_node.getElementType() == PyTokenTypes.DEF_KEYWORD) {
+            PsiElement maybe_func = token.getParent();
+            if (maybe_func instanceof PyFunction) {
+              PyFunction func = (PyFunction)maybe_func;
+              PyUtil.MethodFlags flags = PyUtil.MethodFlags.of(func);
+              if (flags != null) {
+                // we're in a method
+                // TODO: all string constants go to Settings
+                String pname = flags.isClassMethod() || flags.isMetaclassMethod() ? "cls" : "self";
+                final boolean is_new = PyNames.NEW.equals(func.getName());
+                if (flags.isMetaclassMethod() && is_new) pname = "typ";
+                else if (flags.isClassMethod() || is_new) pname = "cls";
+                else if (flags.isStaticMethod()) pname="";
+                documentManager.commitDocument(document);
+                // TODO: only print the ")" if Settings require it
+                EditorModificationUtil.typeInStringAtCaretHonorBlockSelection(editor, "("+pname+"):", true);
+                editor.getCaretModel().moveToOffset(offset + 1 + pname.length()); // right after param name
+                return Result.STOP;
+              }
+            }
+          }
+        }
+      }
+    }
+    return Result.CONTINUE; // the default
+  }
+}

python/src/com/jetbrains/python/inspections/PyMethodParametersInspection.java

@@ -11,16 +11,10 @@ import com.jetbrains.python.PyNames;
 import com.jetbrains.python.actions.AddSelfQuickFix;
 import com.jetbrains.python.actions.RenameParameterQuickFix;
 import com.jetbrains.python.psi.*;
-import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
 
-import java.util.Set;
-
-import static com.jetbrains.python.psi.PyFunction.Flag.CLASSMETHOD;
-import static com.jetbrains.python.psi.PyFunction.Flag.STATICMETHOD;
-
 /**
  * Looks for the 'self' or its equivalents.
  * @author dcheryasov
@@ -50,34 +44,16 @@ public class PyMethodParametersInspection extends PyInspection {
     }
 
 
-    private static boolean among(@NotNull String what, String... variants) {
-      for (String s : variants) {
-        if (what.equals(s)) return true;
-      }
-      return false;
-    }
-
     @Override
     public void visitPyFunction(final PyFunction node) {
-      PsiElement cap = PyUtil.getConcealingParent(node);
-      if (cap instanceof PyClass) {
+      PyUtil.MethodFlags flags = PyUtil.MethodFlags.of(node);
+      if (flags != null) {
         PyParameterList plist = node.getParameterList();
         PyParameter[] params = plist.getParameters();
-        Set<PyFunction.Flag> flags = PyUtil.detectDecorationsAndWrappersOf(node);
-        boolean isMetaclassMethod = false;
-        PyClass type_cls = PyBuiltinCache.getInstance(node).getClass("type");
-        for (PyClass ancestor_cls : ((PyClass)cap).iterateAncestors()) {
-          if (ancestor_cls == type_cls) {
-            isMetaclassMethod = true;
-            break;
-          }
-        }
         final String method_name = node.getName();
-        boolean isSpecialMetaclassMethod = isMetaclassMethod && among(method_name, PyNames.INIT, "__call__");
-        final boolean is_staticmethod = flags.contains(STATICMETHOD);
         if (params.length == 0) {
           // check for "staticmetod"
-          if (is_staticmethod) return; // no params may be fine
+          if (flags.isStaticMethod()) return; // no params may be fine
           // check actual param list
           ASTNode name_node = node.getNameNode();
           if (name_node != null) {
@@ -87,7 +63,7 @@ public class PyMethodParametersInspection extends PyInspection {
               open_paren != null && close_paren != null &&
               "(".equals(open_paren.getText()) && ")".equals(close_paren.getText())
             ) {
-              String paramName = flags.contains(CLASSMETHOD) || isMetaclassMethod ? "cls" : "self";
+              String paramName = flags.isClassMethod() || flags.isMetaclassMethod() ? "cls" : "self";
               registerProblem(
                 plist, PyBundle.message("INSP.must.have.first.parameter", paramName),
                 ProblemHighlightType.GENERIC_ERROR, null, new AddSelfQuickFix(paramName)
@@ -101,20 +77,18 @@ public class PyMethodParametersInspection extends PyInspection {
             String pname = first_param.getText();
             // every dup, swap, drop, or dup+drop of "self"
             @NonNls String[] mangled = {"eslf", "sself", "elf", "felf", "slef", "seelf", "slf", "sslf", "sefl", "sellf", "sef", "seef"};
-            for (String typo : mangled) {
-              if (typo.equals(pname)) {
-                registerProblem(
-                  PyUtil.sure(params[0].getNode()).getPsi(),
-                  PyBundle.message("INSP.probably.mistyped.self"),
-                  new RenameParameterQuickFix(PyNames.CANONICAL_SELF)
-                );
-                return;
-              }
+            if (PyUtil.among(pname, mangled)) {
+              registerProblem(
+                PyUtil.sure(params[0].getNode()).getPsi(),
+                PyBundle.message("INSP.probably.mistyped.self"),
+                new RenameParameterQuickFix(PyNames.CANONICAL_SELF)
+              );
+              return;
             }
             String CLS = "cls"; // TODO: move to style settings
-            if (isMetaclassMethod && PyNames.NEW.equals(method_name)) {
+            if (flags.isMetaclassMethod() && PyNames.NEW.equals(method_name)) {
               final String[] POSSIBLE_PARAM_NAMES = {"typ", "meta"}; // TODO: move to style settings
-              if (!among(pname, POSSIBLE_PARAM_NAMES)) {
+              if (!PyUtil.among(pname, POSSIBLE_PARAM_NAMES)) {
                 registerProblem(
                   PyUtil.sure(params[0].getNode()).getPsi(),
                   PyBundle.message("INSP.usually.named.$0", POSSIBLE_PARAM_NAMES[0]),
@@ -122,7 +96,7 @@ public class PyMethodParametersInspection extends PyInspection {
                 );
               }
             }
-            else if (flags.contains(CLASSMETHOD) || isSpecialMetaclassMethod || PyNames.NEW.equals(method_name)) {
+            else if (flags.isClassMethod() || flags.isSpecialMetaclassMethod() || PyNames.NEW.equals(method_name)) {
               if (!CLS.equals(pname)) {
                 registerProblem(
                   PyUtil.sure(params[0].getNode()).getPsi(),
@@ -131,8 +105,8 @@ public class PyMethodParametersInspection extends PyInspection {
                 );
               }
             }
-            else if (!is_staticmethod && !first_param.isPositionalContainer() && !PyNames.CANONICAL_SELF.equals(pname)) {
-              if (isMetaclassMethod && CLS.equals(pname)) {
+            else if (!flags.isStaticMethod() && !first_param.isPositionalContainer() && !PyNames.CANONICAL_SELF.equals(pname)) {
+              if (flags.isMetaclassMethod() && CLS.equals(pname)) {
                 return;   // accept either 'self' or 'cls' for all methods in metaclass
               }
               registerProblem(
@@ -143,7 +117,7 @@ public class PyMethodParametersInspection extends PyInspection {
             }
           }
           else { // the unusual case of a method with first tuple param
-            if (!is_staticmethod) {
+            if (!flags.isStaticMethod()) {
               registerProblem(plist, PyBundle.message("INSP.first.param.must.not.be.tuple"));
             }
           }

python/src/com/jetbrains/python/psi/PyUtil.java

@@ -678,6 +678,18 @@ public class PyUtil {
     return null;
   }
 
+  /**
+   * @param what thing to search for
+   * @param variants things to search among
+   * @return true iff what.equals() one of the variants.
+   */
+  public static <T> boolean among(@NotNull T what, T... variants) {
+    for (T s : variants) {
+      if (what.equals(s)) return true;
+    }
+    return false;
+  }
+
   public static class UnderscoreFilter implements Condition<String> {
     private int myAllowed; // how many starting underscores is allowed: 0 is none, 1 is only one, 2 is two and more.
 
@@ -716,5 +728,66 @@ public class PyUtil {
     }
     return false;
   }
+
+  public static class MethodFlags {
+
+    private boolean myIsStaticMethod;
+    private boolean myIsMetaclassMethod;
+    private boolean myIsSpecialMetaclassMethod;
+    private boolean myIsClassMethod;
+
+    /**
+     * @return true iff the method belongs to a metaclass (an ancestor of 'type').
+     */
+    public boolean isMetaclassMethod() {
+      return myIsMetaclassMethod;
+    }
+
+    /**
+     * @return iff isMetaclassMethod and the method is either __init__ or __call__.
+     */
+    public boolean isSpecialMetaclassMethod() {
+      return myIsSpecialMetaclassMethod;
+    }
+
+    public boolean isStaticMethod() {
+      return myIsStaticMethod;
+    }
+
+    public boolean isClassMethod() {
+      return myIsClassMethod;
+    }
+
+    private MethodFlags(boolean isClassMethod, boolean isStaticMethod, boolean isMetaclassMethod, boolean isSpecialMetaclassMethod) {
+      myIsClassMethod = isClassMethod;
+      myIsStaticMethod = isStaticMethod;
+      myIsMetaclassMethod = isMetaclassMethod;
+      myIsSpecialMetaclassMethod = isSpecialMetaclassMethod;
+    }
+
+    /**
+     * @param node a function
+     * @return a new flags object, or null if the function is not a method
+     */
+    @Nullable
+    public static MethodFlags of(@NotNull PyFunction node) {
+      PyClass cls = node.getContainingClass();
+      if (cls != null) {
+        Set<PyFunction.Flag> flags = detectDecorationsAndWrappersOf(node);
+        boolean isMetaclassMethod = false;
+        PyClass type_cls = PyBuiltinCache.getInstance(node).getClass("type");
+        for (PyClass ancestor_cls : cls.iterateAncestors()) {
+          if (ancestor_cls == type_cls) {
+            isMetaclassMethod = true;
+            break;
+          }
+        }
+        final String method_name = node.getName();
+        boolean isSpecialMetaclassMethod = isMetaclassMethod && method_name != null && among(method_name, PyNames.INIT, "__call__");
+        return new MethodFlags(flags.contains(CLASSMETHOD), flags.contains(STATICMETHOD), isMetaclassMethod, isSpecialMetaclassMethod);
+      }
+      return null;
+    }
+  }
 }
 

python/testData/editing/firstParamClassmethod.after.py

@@ -0,0 +1,3 @@
+class A(object):
+  @classmethod
+  def foo(cls):

python/testData/editing/firstParamClassmethod.py

@@ -0,0 +1,3 @@
+class A(object):
+  @classmethod
+  def foo<caret>

python/testData/editing/firstParamMetaClass.after.py

@@ -0,0 +1,3 @@
+class A(type):
+  @classmethod
+  def f(cls):

python/testData/editing/firstParamMetaClass.py

@@ -0,0 +1,3 @@
+class A(type):
+  @classmethod
+  def f<caret>

python/testData/editing/firstParamMetaNew.after.py

@@ -0,0 +1,2 @@
+class A(type):
+  def __new__(typ):

python/testData/editing/firstParamMetaNew.py

@@ -0,0 +1,2 @@
+class A(type):
+  def __new__<caret>

python/testData/editing/firstParamMetaSimple.after.py

@@ -0,0 +1,2 @@
+class A(type):
+  def f(cls):

python/testData/editing/firstParamMetaSimple.py

@@ -0,0 +1,2 @@
+class A(type):
+  def f<caret>

python/testData/editing/firstParamSimple.after.py

@@ -0,0 +1,2 @@
+class A(object):
+  def foo(self):

python/testData/editing/firstParamSimple.py

@@ -0,0 +1,2 @@
+class A(object):
+  def foo<caret>

python/testData/editing/firstParamSimpleInit.after.py

@@ -0,0 +1,2 @@
+class A(object):
+  def __init__(self):

python/testData/editing/firstParamSimpleInit.py

@@ -0,0 +1,2 @@
+class A(object):
+  def __init__<caret>

python/testData/editing/firstParamSimpleNew.after.py

@@ -0,0 +1,2 @@
+class A(object):
+  def __new__(cls):

python/testData/editing/firstParamSimpleNew.py

@@ -0,0 +1,2 @@
+class A(object):
+  def __new__<caret>

python/testData/editing/firstParamStaticmethod.after.py

@@ -0,0 +1,3 @@
+class A(object):
+  @staticmethod
+  def foo():

python/testData/editing/firstParamStaticmethod.py

@@ -0,0 +1,3 @@
+class A(object):
+  @staticmethod
+  def foo<caret>

python/testSrc/com/jetbrains/python/PyEditingTest.java

@@ -176,4 +176,55 @@ public class PyEditingTest extends PyLightFixtureTestCase {
     });
     return myFixture.getDocument(file).getText();
   }
+
+  private void doTypingTest(final char character) {
+    final String testName = "editing/" + getTestName(true);
+    myFixture.configureByFile(testName + ".py");
+    doTyping(character);
+    myFixture.checkResultByFile(testName + ".after.py");
+  }
+
+  private void doTyping(final char character) {
+    final int offset = myFixture.getEditor().getCaretModel().getOffset();
+    final PsiFile file = ApplicationManager.getApplication().runWriteAction(new Computable<PsiFile>() {
+      @Override
+      public PsiFile compute() {
+        myFixture.getEditor().getCaretModel().moveToOffset(offset);
+        myFixture.type(character);
+        return myFixture.getFile();
+      }
+    });
+  }
+
+  public void testFirstParamClassmethod() {
+    doTypingTest('(');
+  }
+
+  public void testFirstParamMetaClass() {
+    doTypingTest('(');
+  }
+
+  public void testFirstParamMetaNew() {
+    doTypingTest('(');
+  }
+
+  public void testFirstParamMetaSimple() {
+    doTypingTest('(');
+  }
+
+  public void testFirstParamSimpleInit() {
+    doTypingTest('(');
+  }
+
+  public void testFirstParamSimpleNew() {
+    doTypingTest('(');
+  }
+
+  public void testFirstParamSimple() {
+    doTypingTest('(');
+  }
+
+  public void testFirstParamStaticmethod() {
+    doTypingTest('(');
+  }
 }