fixed PY-9654 Make function from method/Make method static: correctly update class calls with first instance argument

python/src/com/jetbrains/python/inspections/quickfix/PyMakeFunctionFromMethodQuickFix.java

@@ -5,6 +5,7 @@ import com.intellij.codeInspection.ProblemDescriptor;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiReference;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.usageView.UsageInfo;
 import com.jetbrains.python.PyBundle;
@@ -55,8 +56,64 @@ public class PyMakeFunctionFromMethodQuickFix implements LocalQuickFix {
     for (UsageInfo usage : usages) {
       final PsiElement usageElement = usage.getElement();
       if (usageElement instanceof PyReferenceExpression) {
-        PyUtil.removeQualifier((PyReferenceExpression)usageElement);
+        updateUsage((PyReferenceExpression)usageElement);
       }
     }
   }
+
+  private static void updateUsage(@NotNull final PyReferenceExpression element) {
+    final PyExpression qualifier = element.getQualifier();
+    if (qualifier == null) return;
+
+    if (qualifier instanceof PyCallExpression) {              // remove qualifier foo.A().m()
+      PyUtil.removeQualifier(element);
+    }
+    else {
+      final PsiReference reference = qualifier.getReference();
+      if (reference == null) return;
+
+      final PsiElement resolved = reference.resolve();
+      if (resolved instanceof PyTargetExpression) {  // qualifier came from assignment  a = A(); a.m()
+        updateAssignment(element, resolved);
+      }
+      else if (resolved instanceof PyClass) {     //call with first instance argument A.m(A())
+        final PsiElement dot = qualifier.getNextSibling();
+        if (dot != null) dot.delete();
+        qualifier.delete();
+        updateArgumentList(element);
+      }
+    }
+  }
+
+  private static void updateAssignment(PyReferenceExpression element, @NotNull final PsiElement resolved) {
+    final PsiElement parent = resolved.getParent();
+    if (parent instanceof PyAssignmentStatement) {
+      final PyExpression value = ((PyAssignmentStatement)parent).getAssignedValue();
+      if (value instanceof PyCallExpression) {
+        final PyExpression callee = ((PyCallExpression)value).getCallee();
+        if (callee instanceof PyReferenceExpression) {
+          final PyExpression calleeQualifier = ((PyReferenceExpression)callee).getQualifier();
+          if (calleeQualifier != null) {
+            value.replace(calleeQualifier);
+          }
+          else {
+            PyUtil.removeQualifier(element);
+          }
+        }
+      }
+    }
+  }
+
+  private static void updateArgumentList(@NotNull final PyReferenceExpression element) {
+    final PyCallExpression callExpression = PsiTreeUtil.getParentOfType(element, PyCallExpression.class);
+    if (callExpression == null) return;
+    PyArgumentList argumentList = callExpression.getArgumentList();
+    if (argumentList == null) return;
+    final PyExpression[] arguments = argumentList.getArguments();
+    if (arguments.length > 0) {
+      final PyExpression argument = arguments[0];
+      PyUtil.eraseWhitespaceAndComma(argument.getParent().getNode(), argument, false);
+      argument.delete();
+    }
+  }
 }

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

@@ -342,7 +342,7 @@ public class PyUtil {
    * @param backwards   true to erase prev siblings, false to erase next siblings.
    * @return true       if a comma was found and removed.
    */
-  private static boolean eraseWhitespaceAndComma(ASTNode parent_node, PsiElement item, boolean backwards) {
+  public static boolean eraseWhitespaceAndComma(ASTNode parent_node, PsiElement item, boolean backwards) {
     // we operate on AST, PSI won't let us delete whitespace easily.
     boolean is_comma;
     boolean got_comma = false;
@@ -673,25 +673,9 @@ public class PyUtil {
       element.replace(expression);
     }
     else {
-      final PsiReference reference = qualifier.getReference();
-      if (reference != null) {
-        final PsiElement resolved = reference.resolve();
-        if (resolved instanceof PyTargetExpression) {
-          final PsiElement parent = resolved.getParent();
-          if (parent instanceof PyAssignmentStatement) {
-            final PyExpression value = ((PyAssignmentStatement)parent).getAssignedValue();
-            if (value instanceof PyCallExpression) {
-              final PyExpression callee = ((PyCallExpression)value).getCallee();
-              if (callee instanceof PyReferenceExpression) {
-                final PyExpression calleeQualifier = ((PyReferenceExpression)callee).getQualifier();
-                if (calleeQualifier != null) {
-                  value.replace(calleeQualifier);
-                }
-              }
-            }
-          }
-        }
-      }
+      final PsiElement dot = qualifier.getNextSibling();
+      if (dot != null) dot.delete();
+      qualifier.delete();
     }
   }
 

python/testData/quickFixes/PyMakeFunctionFromMethodQuickFixTest/test.py

@@ -0,0 +1,6 @@
+__author__ = 'ktisha'
+
+class A:
+    def fo<caret>o(self):
+        print("Hello Pycharm!")
+

python/testData/quickFixes/PyMakeFunctionFromMethodQuickFixTest/usageAssignment.py

@@ -0,0 +1,6 @@
+class A:
+    def m<caret>(self):
+        print 1
+
+b = A()
+a = b.m()

python/testData/quickFixes/PyMakeFunctionFromMethodQuickFixTest/usageAssignment_after.py

@@ -0,0 +1,9 @@
+class A: pass
+
+
+def m():
+    print 1
+
+
+b = A()
+a = m()

python/testData/quickFixes/PyMakeFunctionFromMethodQuickFixTest/usageClassCallArgument.py

@@ -0,0 +1,5 @@
+class A:
+    def m<caret>(self, x):
+        print 1
+
+A.m(A(), 1)

python/testData/quickFixes/PyMakeFunctionFromMethodQuickFixTest/usageClassCallArgument_after.py

@@ -0,0 +1,8 @@
+class A: pass
+
+
+def m(x):
+    print 1
+
+
+m(1)

python/testData/quickFixes/PyMakeFunctionFromMethodQuickFixTest/usageImport.py

@@ -0,0 +1,5 @@
+import test
+
+b = test.A()
+
+b.foo(1)

python/testData/quickFixes/PyMakeFunctionFromMethodQuickFixTest/usageImport1.py

@@ -0,0 +1,3 @@
+import test
+
+test.A().foo(1)

python/testData/quickFixes/PyMakeFunctionFromMethodQuickFixTest/usageImport1_after.py

@@ -0,0 +1,3 @@
+import test
+
+test.foo(1)

python/testData/quickFixes/PyMakeFunctionFromMethodQuickFixTest/usageImport_after.py

@@ -0,0 +1,5 @@
+import test
+
+b = test
+
+b.foo(1)

python/testSrc/com/jetbrains/python/quickFixes/PyMakeFunctionFromMethodQuickFixTest.java

@@ -1,5 +1,6 @@
 package com.jetbrains.python.quickFixes;
 
+import com.intellij.codeInsight.intention.IntentionAction;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.inspections.PyMethodMayBeStaticInspection;
 
@@ -36,4 +37,30 @@ public class PyMakeFunctionFromMethodQuickFixTest extends PyQuickFixTestCase {
     doInspectionTest(PyMethodMayBeStaticInspection.class, PyBundle.message("QFIX.NAME.make.function"));
   }
 
+  public void testUsageClassCallArgument() {
+    doInspectionTest(PyMethodMayBeStaticInspection.class, PyBundle.message("QFIX.NAME.make.function"));
+  }
+
+  public void testUsageAssignment() {
+    doInspectionTest(PyMethodMayBeStaticInspection.class, PyBundle.message("QFIX.NAME.make.function"));
+  }
+
+  public void testUsageImport() {
+    doTest();
+  }
+
+  public void testUsageImport1() {
+    doTest();
+  }
+
+
+  protected void doTest() {
+    final String testFileName = getTestName(true);
+    myFixture.enableInspections(PyMethodMayBeStaticInspection.class);
+    myFixture.configureByFiles("test.py", testFileName + ".py");
+    final IntentionAction intentionAction = myFixture.findSingleIntention(PyBundle.message("QFIX.NAME.make.function"));
+    assertNotNull(intentionAction);
+    myFixture.launchAction(intentionAction);
+    myFixture.checkResultByFile(testFileName + ".py", testFileName + "_after.py", true);
+  }
 }