'Override method' inserts super method call instead of 'pass' (PY-302)

python/src/com/jetbrains/python/PyNames.java

@@ -20,6 +20,9 @@ public class PyNames {
   @NonNls public static final String GETATTR = "__getattr__";
   @NonNls public static final String GETATTRIBUTE = "__getattribute__";
   @NonNls public static final String CLASS = "__class__";
+  @NonNls public static final String METACLASS = "__metaclass__";
+
+  @NonNls public static final String SUPER = "super";
 
   @NonNls public static final String OBJECT = "object";
   @NonNls public static final String NONE = "None";

python/src/com/jetbrains/python/codeInsight/override/PyOverrideImplementUtil.java

@@ -16,8 +16,9 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiWhiteSpace;
 import com.intellij.psi.util.PsiTreeUtil;
-import com.intellij.util.IncorrectOperationException;
+import com.jetbrains.python.PyNames;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyFunctionBuilder;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -30,6 +31,9 @@ import java.util.*;
 public class PyOverrideImplementUtil {
   private static final Logger LOG = Logger.getInstance("#com.jetbrains.python.codeInsight.override.PyOverrideImplementUtil");
 
+  private PyOverrideImplementUtil() {
+  }
+
   @Nullable
   public static PyClass getContextClass(@NotNull final Project project, @NotNull final Editor editor, @NotNull final PsiFile file) {
     PsiDocumentManager.getInstance(project).commitAllDocuments();
@@ -90,22 +94,19 @@ public class PyOverrideImplementUtil {
     overrideMethods(editor, pyClass, membersToOverride);
   }
 
-  public static void overrideMethods(final Editor editor, final PyClass pyClass, List<PyMethodMember> membersToOverride) {
-    final List<String> newMembers = generateCode(membersToOverride);
-    if (newMembers.isEmpty()) {
+  public static void overrideMethods(final Editor editor, final PyClass pyClass, final List<PyMethodMember> membersToOverride) {
+    if (membersToOverride == null) {
       return;
     }
-
     new WriteCommandAction(pyClass.getProject(), pyClass.getContainingFile()) {
       protected void run(final Result result) throws Throwable {
-        write(pyClass, newMembers, pyClass.getProject(), editor);
+        write(pyClass, membersToOverride, editor);
       }
     }.execute();
   }
 
   private static void write(@NotNull final PyClass pyClass,
-                            @NotNull final List<String> newMembers,
-                            @NotNull final Project project,
+                            @NotNull final List<PyMethodMember> newMembers,
                             @NotNull final Editor editor) {
     final PyStatementList statementList = pyClass.getStatementList();
     final int offset = editor.getCaretModel().getOffset();
@@ -117,15 +118,11 @@ public class PyOverrideImplementUtil {
     }
 
     PyFunction element = null;
-    for (String newMember : newMembers) {
-      element = PyElementGenerator.getInstance(project).createFromText(PyFunction.class, newMember + "\n    pass");
-      try {
-        element = (PyFunction)statementList.addAfter(element, anchor);
-        element = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(element);
-      }
-      catch (IncorrectOperationException e) {
-        LOG.error(e);
-      }
+    for (PyMethodMember newMember : newMembers) {
+      PyFunction baseFunction = (PyFunction) newMember.getPsiElement();
+      final PyFunctionBuilder builder = buildOverriddenFunction(pyClass, baseFunction);
+      PyFunction function = builder.addFunctionAfter(statementList, anchor);
+      element = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(function);
     }
     
     PyPsiUtils.removeRedundantPass(statementList);
@@ -135,31 +132,43 @@ public class PyOverrideImplementUtil {
     editor.getSelectionModel().setSelection(start, element.getTextRange().getEndOffset());
   }
 
-  private static List<String> generateCode(final List<PyMethodMember> members) {
-    if (members == null) {
-      return Collections.emptyList();
+  private static PyFunctionBuilder buildOverriddenFunction(PyClass pyClass, PyFunction baseFunction) {
+    PyFunctionBuilder pyFunctionBuilder = new PyFunctionBuilder(baseFunction.getName());
+    final PyDecoratorList decorators = baseFunction.getDecoratorList();
+    if (decorators != null && decorators.findDecorator(PyNames.CLASSMETHOD) != null) {
+      pyFunctionBuilder.decorate(PyNames.CLASSMETHOD);
     }
-    List<String> newMembers = new ArrayList<String>();
-    for (PyMethodMember member : members) {
-      newMembers.add(generateNewMethod(member.getPsiElement()));
+    final PyParameter[] baseParams = baseFunction.getParameterList().getParameters();
+    for (PyParameter parameter : baseParams) {
+      pyFunctionBuilder.parameter(parameter.getText());
     }
-    return newMembers;
-  }
 
-  @NotNull
-  private static String generateNewMethod(@NotNull final PsiElement element) {
-    assert (element instanceof PyFunction);
-    final PyFunction function = (PyFunction)element;
-    final StringBuilder newMethodText = new StringBuilder();
-    final PyDecoratorList decoratorList = function.getDecoratorList();
-    if (decoratorList != null) {
-      for (PyDecorator decorator: decoratorList.getDecorators()) {
-        if ("classmethod".equals(decorator.getCallee().getText())) {
-          newMethodText.append("@classmethod\n");
-        }
+    PyClass baseClass = baseFunction.getContainingClass();
+    assert baseClass != null;
+    StringBuilder statementBody = new StringBuilder();
+    if (baseClass.isNewStyleClass()) {
+      statementBody.append(PyNames.SUPER);
+      statementBody.append("(");
+      final LanguageLevel langLevel = ((PyFile)pyClass.getContainingFile()).getLanguageLevel();
+      if (!langLevel.isPy3K()) {
+        statementBody.append(pyClass.getName()).append(", self");
       }
+      statementBody.append(").").append(baseFunction.getName()).append("(");
+      for (int i = 1; i < baseParams.length; i++) {
+        statementBody.append(baseParams [i].getText());
+      }
+      statementBody.append(")");
     }
-    return newMethodText.append("def ").append(function.getName()).append(function.getParameterList().getText()).append(":").toString();
+    else {
+      statementBody.append(baseClass.getName()).append(".").append(baseFunction.getName()).append("(");
+      for (PyParameter param : baseParams) {
+        statementBody.append(param.getText());
+      }
+      statementBody.append(")");
+    }
+
+    pyFunctionBuilder.statement(statementBody.toString());
+    return pyFunctionBuilder;
   }
 
   @NotNull
@@ -173,7 +182,4 @@ public class PyOverrideImplementUtil {
     }
     return superFunctions.values();
   }
-
-  private PyOverrideImplementUtil() {
-  }
 }

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

@@ -1,10 +1,10 @@
 package com.jetbrains.python.inspections;
 
-import com.intellij.codeInspection.LocalInspectionTool;
 import com.intellij.codeInspection.ProblemsHolder;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiElementVisitor;
 import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.PyNames;
 import com.jetbrains.python.psi.PyCallExpression;
 import com.jetbrains.python.psi.PyClass;
 import com.jetbrains.python.psi.PyExpression;
@@ -37,7 +37,7 @@ public class PySuperArgumentsInspection extends PyInspection {
 
     @Override
     public void visitPyCallExpression(PyCallExpression node) {
-      if ("super".equals(node.getCallee().getName())) {
+      if (PyNames.SUPER.equals(node.getCallee().getName())) {
         PyExpression[] arguments = node.getArguments();
         if (arguments.length == 2) {
           if (arguments[0] instanceof PyReferenceExpression && arguments[1] instanceof PyReferenceExpression) {

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

@@ -2,19 +2,20 @@ package com.jetbrains.python.psi;
 
 import com.intellij.psi.StubBasedPsiElement;
 import com.jetbrains.python.psi.stubs.PyDecoratorListStub;
-import com.jetbrains.python.psi.PyDecorator;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 /**
  * A list of function decorators.
- * User: dcheryasov
- * Date: Sep 28, 2008
+ * @author dcheryasov
  */
 public interface PyDecoratorList extends PyElement, StubBasedPsiElement<PyDecoratorListStub> {
-
   /**
    * @return decorators of function, in order of declaration (outermost first).
    */
   @NotNull
   PyDecorator[] getDecorators();
+
+  @Nullable
+  PyDecorator findDecorator(String name);
 }

python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java

@@ -72,7 +72,7 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress
       PyExpression callee = getCallee();
       if (callee instanceof PyReferenceExpression) {
         // hardwired special cases
-        if ("super".equals(callee.getText())) {
+        if (PyNames.SUPER.equals(callee.getText())) {
           final PyType superCallType = getSuperCallType(callee, context);
           if (superCallType != null) {
             return superCallType;
@@ -135,7 +135,7 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress
     PsiElement must_be_super_init = ((PyReferenceExpression)callee).getReference().resolve();
     if (must_be_super_init instanceof PyFunction) {
       PyClass must_be_super = ((PyFunction)must_be_super_init).getContainingClass();
-      if (must_be_super == PyBuiltinCache.getInstance(this).getClass("super")) {
+      if (must_be_super == PyBuiltinCache.getInstance(this).getClass(PyNames.SUPER)) {
         PyArgumentList arglist = getArgumentList();
         if (arglist != null) {
           PyExpression[] args = arglist.getArguments();

python/src/com/jetbrains/python/psi/impl/PyClassImpl.java

@@ -708,6 +708,9 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
   }
 
   private boolean calculateNewStyleClass() {
+    if (((PyFile) getContainingFile()).getLanguageLevel().isPy3K()) {
+      return true;
+    }
     PyClass objclass = PyBuiltinCache.getInstance(this).getClass("object");
     if (this == objclass) return true; // a rare but possible case
     if (hasNewStyleMetaClass(this)) return true;
@@ -723,7 +726,7 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
   private static boolean hasNewStyleMetaClass(PyClass pyClass) {
     final PsiFile containingFile = pyClass.getContainingFile();
     if (containingFile instanceof PyFile) {
-      final PsiElement element = ((PyFile)containingFile).findExportedName("__metaclass__");
+      final PsiElement element = ((PyFile)containingFile).findExportedName(PyNames.METACLASS);
       if (element instanceof PyTargetExpression) {
         final PyExpression assignedValue = ((PyTargetExpression)element).findAssignedValue();
         if (assignedValue != null && assignedValue.getText().equals("type")) {

python/src/com/jetbrains/python/psi/impl/PyDecoratorListImpl.java

@@ -32,4 +32,15 @@ public class PyDecoratorListImpl extends PyBaseElementImpl<PyDecoratorListStub>
     return getStubOrPsiChildren(PyElementTypes.DECORATOR_CALL, decoarray);
     //return decoarray;
   }
+
+  @Override
+  public PyDecorator findDecorator(String name) {
+    final PyDecorator[] decorators = getDecorators();
+    for (PyDecorator decorator : decorators) {
+      if (name.equals(decorator.getCallee().getText())) {
+        return decorator;
+      }
+    }
+    return null;
+  }
 }

python/src/com/jetbrains/python/psi/impl/PyFunctionBuilder.java

@@ -6,7 +6,6 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.codeStyle.CodeStyleSettings;
 import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
 import com.jetbrains.python.PythonFileType;
-import com.jetbrains.python.PythonLanguage;
 import com.jetbrains.python.psi.PyElementGenerator;
 import com.jetbrains.python.psi.PyFunction;
 
@@ -21,6 +20,7 @@ public class PyFunctionBuilder {
   private final String myName;
   private final List<String> myParameters = new ArrayList<String>();
   private final List<String> myStatements = new ArrayList<String>();
+  private final List<String> myDecorators = new ArrayList<String>();
 
   public PyFunctionBuilder(String name) {
     myName = name;
@@ -46,6 +46,10 @@ public class PyFunctionBuilder {
     return (PyFunction) target.add(buildFunction(target.getProject()));
   }
 
+  public PyFunction addFunctionAfter(PsiElement target, PsiElement anchor) {
+    return (PyFunction) target.addAfter(buildFunction(target.getProject()), anchor);
+  }
+
   public PyFunction buildFunction(Project project) {
     String text = buildText(project);
     PyElementGenerator generator = PyElementGenerator.getInstance(project);
@@ -53,7 +57,11 @@ public class PyFunctionBuilder {
   }
 
   private String buildText(Project project) {
-    StringBuilder builder = new StringBuilder("def ");
+    StringBuilder builder = new StringBuilder();
+    for (String decorator : myDecorators) {
+      builder.append(decorator).append("\n");
+    }
+    builder.append("def ");
     builder.append(myName).append("(");
     builder.append(StringUtil.join(myParameters, ", "));
     builder.append("):");
@@ -66,4 +74,8 @@ public class PyFunctionBuilder {
     }
     return builder.toString();
   }
+
+  public void decorate(String decoratorName) {
+    myDecorators.add("@" + decoratorName);
+  }
 }

python/src/com/jetbrains/python/validation/UnsupportedFeatures.java

@@ -5,6 +5,7 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiWhiteSpace;
 import com.intellij.util.containers.HashSet;
 import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.PyNames;
 import com.jetbrains.python.codeInsight.intentions.*;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyQualifiedName;
@@ -133,7 +134,7 @@ public class UnsupportedFeatures extends PyAnnotator {
       final PsiElement firstChild = node.getFirstChild();
       if (firstChild != null) {
         final String name = firstChild.getText();
-        if ("super".equals(name)) {
+        if (PyNames.SUPER.equals(name)) {
           final PyArgumentList argumentList = node.getArgumentList();
           if (argumentList != null && argumentList.getArguments().length == 0) {
             getHolder().createWarningAnnotation(node, "super() should have arguments in Python 2");

python/testData/override/classmethod_after.py

@@ -6,4 +6,4 @@ class A:
 class B(A):
     @classmethod
     def foo(cls):
-        <selection>pass</selection>
+        <selection>A.foo(cls)</selection>

python/testData/override/newStyle.py

@@ -0,0 +1,6 @@
+class A(object):
+    def m(self):
+        pass
+
+class B(A):
+    <caret>pass

python/testData/override/newStyle_after.py

@@ -0,0 +1,8 @@
+class A(object):
+    def m(self):
+        pass
+
+class B(A):
+    def m(self):
+        <selection>super(B, self).m()</selection>
+

python/testData/override/py3k.py

@@ -0,0 +1,6 @@
+class A:
+    def m(self):
+        pass
+
+class B(A):
+    <caret>pass

python/testData/override/py3k_after.py

@@ -0,0 +1,8 @@
+class A:
+    def m(self):
+        pass
+
+class B(A):
+    def m(self):
+        <selection>super().m()</selection>
+

python/testData/override/simple_after.py

@@ -3,7 +3,7 @@ class A:
 
 class B(A):
     def doStuff(self):
-        <selection>pass</selection>
+        <selection>A.doStuff(self)</selection>
 
     def otherMethod(self, foo, bar):
         print foo, bar

python/testSrc/com/jetbrains/python/PyOverrideTest.java

@@ -3,9 +3,11 @@ package com.jetbrains.python;
 import com.jetbrains.python.codeInsight.override.PyMethodMember;
 import com.jetbrains.python.codeInsight.override.PyOverrideImplementUtil;
 import com.jetbrains.python.fixtures.PyLightFixtureTestCase;
+import com.jetbrains.python.psi.LanguageLevel;
 import com.jetbrains.python.psi.PyClass;
 import com.jetbrains.python.psi.PyFile;
 import com.jetbrains.python.psi.PyFunction;
+import com.jetbrains.python.psi.impl.PythonLanguageLevelPusher;
 
 import java.util.Collections;
 import java.util.List;
@@ -14,21 +16,35 @@ import java.util.List;
  * @author yole
  */
 public class PyOverrideTest extends PyLightFixtureTestCase {
-  private void doTest() throws Exception {
+  private void doTest() {
     myFixture.configureByFile("override/" + getTestName(true) + ".py");
     PyFile file = (PyFile) myFixture.getFile();
     List<PyClass> classes = file.getTopLevelClasses();
     PyFunction toOverride = classes.get(0).getMethods() [0];
     PyOverrideImplementUtil.overrideMethods(myFixture.getEditor(), classes.get(1),
                                             Collections.singletonList(new PyMethodMember(toOverride)));
-    myFixture.checkResultByFile("override/" + getTestName(true) + "_after.py");
+    myFixture.checkResultByFile("override/" + getTestName(true) + "_after.py", true);
   }
 
-  public void testSimple() throws Exception {
+  public void testSimple() {
     doTest();
   }
 
-  public void testClassmethod() throws Exception {
+  public void testClassmethod() {
     doTest();
   }
+
+  public void testNewStyle() {
+    doTest();
+  }
+
+  public void testPy3k() {
+    PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), LanguageLevel.PYTHON31);
+    try {
+      doTest();
+    }
+    finally {
+      PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), null);
+    }
+  }
 }