Py-419 and much more

First bunch of tests on extract method.

python/src/com/jetbrains/python/codeInsight/codeFragment/PyCodeFragmentBuilder.java

@@ -30,6 +30,15 @@ public class PyCodeFragmentBuilder extends PyRecursiveElementVisitor {
     endOffset = end;
   }
 
+  @Override
+  public void visitPyAugAssignmentStatement(final PyAugAssignmentStatement node) {
+    final PyExpression target = node.getTarget();
+    if (target instanceof PyReferenceExpression){
+      visitPyReferenceExpression((PyReferenceExpression) target);
+      processDeclaration(target);
+    }
+  }
+
   @Override
   public void visitPyTargetExpression(final PyTargetExpression node) {
     processDeclaration(node);

python/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java

@@ -57,6 +57,16 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
     }
   }
 
+  @Override
+  public void visitPyAugAssignmentStatement(final PyAugAssignmentStatement node) {
+    myBuilder.startNode(node);
+    final PyExpression value = node.getValue();
+    if (value != null){
+      value.accept(this);
+    }
+    node.getTarget().accept(this);
+  }
+
   @Override
   public void visitPyTargetExpression(final PyTargetExpression node) {
     final WriteInstruction instruction = new WriteInstruction(myBuilder, node, node.getName());

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

@@ -17,6 +17,7 @@
 package com.jetbrains.python.psi;
 
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 /**
  * Created by IntelliJ IDEA.
@@ -26,5 +27,8 @@ import org.jetbrains.annotations.NotNull;
  * To change this template use File | Settings | File Templates.
  */
 public interface PyAugAssignmentStatement extends PyStatement {
-    @NotNull PyExpression getTarget();
+  @NotNull
+  PyExpression getTarget();
+  @Nullable
+  PyExpression getValue();
 }

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

@@ -22,6 +22,7 @@ import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.psi.PyAugAssignmentStatement;
 import com.jetbrains.python.psi.PyElementVisitor;
 import com.jetbrains.python.psi.PyExpression;
+import org.jetbrains.annotations.Nullable;
 
 /**
  * Created by IntelliJ IDEA.
@@ -42,10 +43,15 @@ public class PyAugAssignmentStatementImpl extends PyElementImpl implements PyAug
 
   @NotNull
   public PyExpression getTarget() {
-    PyExpression target = childToPsi(PyElementTypes.EXPRESSIONS, 0);
+    final PyExpression target = childToPsi(PyElementTypes.EXPRESSIONS, 0);
     if (target == null) {
       throw new RuntimeException("Target missing in augmented assignment statement");
     }
     return target;
   }
+
+  @Nullable
+  public PyExpression getValue() {
+    return childToPsi(PyElementTypes.EXPRESSIONS, 1);
+  }
 }

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

@@ -13,10 +13,7 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonLanguage;
-import com.jetbrains.python.psi.PyElement;
-import com.jetbrains.python.psi.PyElementType;
-import com.jetbrains.python.psi.PyFile;
-import com.jetbrains.python.psi.PyStatementList;
+import com.jetbrains.python.psi.*;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 

python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodHandler.java

@@ -111,6 +111,7 @@ public class PyExtractMethodHandler implements RefactoringActionHandler {
         return;
       }
       PyExtractMethodUtil.extractFromExpression(project, editor, fragment, expression);
+      return;
     }
 
     CommonRefactoringUtil.showErrorHint(project, editor,

python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java

@@ -9,6 +9,8 @@ import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiNamedElement;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.refactoring.RefactoringBundle;
 import com.intellij.refactoring.RefactoringFactory;
 import com.intellij.refactoring.extractMethod.AbstractExtractMethodDialog;
@@ -21,6 +23,7 @@ import com.jetbrains.python.PythonLanguage;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
@@ -29,6 +32,8 @@ import java.util.Map;
  */
 public class PyExtractMethodUtil {
 
+  public static final String NAME = "extract.method.name";
+
   public static void extractFromStatements(final Project project,
                                            final Editor editor,
                                            final CodeFragment fragment,
@@ -36,7 +41,7 @@ public class PyExtractMethodUtil {
                                            final PsiElement statement2) {
     if (!fragment.getOutputVariables().isEmpty() && fragment.isReturnInstructonInside()) {
       CommonRefactoringUtil.showErrorHint(project, editor,
-                                          "Cannot extract method with non empty output variables and return instructions inside",
+                                          "Cannot perform refactoring from expression with local variables modifications and return instructions inside code fragment",
                                           RefactoringBundle.message("error.title"), "refactoring.extractMethod");
       return;
     }
@@ -126,13 +131,19 @@ public class PyExtractMethodUtil {
                                            final Editor editor,
                                            final CodeFragment fragment,
                                            final PsiElement expression) {
-    if (!fragment.getOutputVariables().isEmpty() && fragment.isReturnInstructonInside()){
+    if (!fragment.getOutputVariables().isEmpty()){
       CommonRefactoringUtil.showErrorHint(project, editor,
-                                          "Cannot extract method with non empty output variables and return instructions inside",
+                                          "Cannot perform refactoring from expression with local variables modifications inside code fragment",
                                           RefactoringBundle.message("error.title"), "refactoring.extractMethod");
       return;
     }
 
+    if (fragment.isReturnInstructonInside()){
+      CommonRefactoringUtil.showErrorHint(project, editor,
+                                          "Cannot extract method with return instructions inside code fragment",
+                                          RefactoringBundle.message("error.title"), "refactoring.extractMethod");
+      return;
+    }
     final Pair<String, AbstractVariableData[]> data = getNameAndVariableData(project, fragment, expression);
     if (data.first == null || data.second == null) {
       return;
@@ -221,14 +232,17 @@ public class PyExtractMethodUtil {
     if (data != null){
       anchor = data.first;
     }
-    // Handle extracting within functions
-    final PsiElement compoundStatement = PyPsiUtils.getCompoundStatement(anchor);
-    final PsiElement parent = compoundStatement.getParent();
-    if (parent instanceof PyFunction){
-      parent.getParent().addBefore(generatedMethod, parent);
+    final PsiNamedElement parent = PsiTreeUtil.getParentOfType(anchor, PyFile.class, PyFunction.class);
+    if (parent instanceof PyFile) {
+      final PsiElement statement = PyPsiUtils.getStatement(parent, anchor);
+      parent.addBefore(generatedMethod, statement);
+      return;
     }
-    final PsiElement statement = PyPsiUtils.getStatement(compoundStatement, anchor);
-    compoundStatement.addBefore(generatedMethod, statement);
+    if (parent instanceof PyFunction) {
+      parent.getParent().addBefore(generatedMethod, parent);
+      return;
+    }
+    throw new IllegalStateException("Compound statement should not be null");
   }
 
   //  Creates string for method parameters
@@ -276,8 +290,25 @@ public class PyExtractMethodUtil {
   }
 
   private static Pair<String, AbstractVariableData[]> getNameAndVariableData(final Project project,
-                                                                     final CodeFragment fragment,
-                                                                     final PsiElement element) {
+                                                                             final CodeFragment fragment,
+                                                                             final PsiElement element) {
+    if (ApplicationManager.getApplication().isUnitTestMode()){
+
+      String name = System.getProperty(NAME);
+      if (name == null){
+        name = "foo";
+      }
+      final List<AbstractVariableData> data = new ArrayList<AbstractVariableData>();
+      for (String in : fragment.getInputVariables()) {
+        final AbstractVariableData d = new AbstractVariableData();
+        d.name = in+"_new";
+        d.originalName = in;
+        d.passAsParameter = true;
+        data.add(d);
+      }
+      return Pair.create(name, data.toArray(new AbstractVariableData[data.size()]));
+    }
+
     final ExtractMethodValidator validator = new ExtractMethodValidator() {
       public String check(final String name) {
         // TODO[oleg] implement context for name clashes

python/testData/codeInsight/controlflow/assignment.py

@@ -0,0 +1 @@
+aaa = aaa + bbb * ccc

python/testData/codeInsight/controlflow/assignment.txt

@@ -0,0 +1,4 @@
+0(1) element: null
+1(2) element: PyAssignmentStatement
+2(3) WRITE ACCESS: aaa
+3() element: null

python/testData/codeInsight/controlflow/augassignment.py

@@ -0,0 +1,2 @@
+aaa = 123
+aaa += bbb * ccc

python/testData/codeInsight/controlflow/augassignment.txt

@@ -0,0 +1,5 @@
+0(1) element: null
+1(2) element: PyAssignmentStatement
+2(3) WRITE ACCESS: aaa
+3(4) element: PyAugAssignmentStatement
+4() element: null

python/testData/refactoring/extractmethod/comment.after.py

@@ -0,0 +1,5 @@
+#Comment to method
+def bar():
+    print("Hello")
+def foo():
+  bar()

python/testData/refactoring/extractmethod/comment.before.py

@@ -0,0 +1,3 @@
+#Comment to method
+def foo():
+  <selection>print("Hello")</selection>

python/testSrc/com/jetbrains/python/PyControlFlowBuilderTest.java

@@ -70,6 +70,14 @@ public class PyControlFlowBuilderTest extends LightMarkedTestCase {
     doTest();
   }
 
+  public void testAssignment() throws Exception {
+    doTest();
+  }
+
+  public void testAugAssignment() throws Exception {
+    doTest();
+  }
+
   public void testFunction() throws Exception {
     final String testName = getTestName(false).toLowerCase();
     configureByFile(testName + ".py");

python/testSrc/com/jetbrains/python/refactoring/FileDataContext.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2000-2009 JetBrains s.r.o.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package com.jetbrains.python.refactoring;
+
+import com.intellij.openapi.actionSystem.DataContext;
+import com.intellij.openapi.actionSystem.LangDataKeys;
+import com.intellij.openapi.actionSystem.PlatformDataKeys;
+import com.intellij.psi.PsiFile;
+import org.jetbrains.annotations.NonNls;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * @author oleg
+*/
+public class FileDataContext implements DataContext {
+  private final PsiFile myFile;
+
+  public FileDataContext(final PsiFile file) {
+    myFile = file;
+  }
+
+  @Nullable
+  public Object getData(@NonNls String dataId) {
+    if (LangDataKeys.LANGUAGE.is(dataId)) {
+      return myFile.getLanguage();
+    }
+    if (PlatformDataKeys.PROJECT.is(dataId)) {
+      return myFile.getProject();
+    }
+    if (LangDataKeys.PSI_FILE.is(dataId)) {
+      return myFile;
+    }
+
+    throw new IllegalArgumentException("Data not supported: " + dataId);
+  }
+}

python/testSrc/com/jetbrains/python/refactoring/PyExtractMethodTest.java

@@ -0,0 +1,106 @@
+package com.jetbrains.python.refactoring;
+
+import com.intellij.lang.LanguageRefactoringSupport;
+import com.intellij.openapi.util.Pair;
+import com.intellij.refactoring.RefactoringActionHandler;
+import com.jetbrains.python.PythonLanguage;
+import com.jetbrains.python.PythonTestUtil;
+import com.jetbrains.python.fixtures.LightMarkedTestCase;
+import com.jetbrains.python.refactoring.extractmethod.PyExtractMethodUtil;
+
+/**
+ * @author oleg
+ */
+public class PyExtractMethodTest extends LightMarkedTestCase {
+  public String getTestDataPath() {
+    return PythonTestUtil.getTestDataPath() + "/refactoring/extractmethod/";
+  }
+
+  private void doTest(final String testPath,
+                      final String name,
+                      final String result,
+                      final Pair<String, String>... files2Create) throws Exception {
+
+    // Create additional files
+    for (Pair<String, String> pair : files2Create) {
+      myFixture.addFileToProject(pair.first, pair.second);
+    }
+    myFixture.configureByFile(testPath);
+    final RefactoringActionHandler handler = LanguageRefactoringSupport.INSTANCE.forLanguage(PythonLanguage.getInstance()).getExtractMethodHandler();
+    try {
+      System.setProperty(PyExtractMethodUtil.NAME, name);
+      handler.invoke(myFixture.getProject(), myFixture.getEditor(), myFixture.getFile(), new FileDataContext(myFixture.getFile()));
+    } catch (Exception e) {
+      assertEquals(result, e.getMessage().trim());
+      return;
+    } finally {
+      System.clearProperty(PyExtractMethodUtil.NAME);
+    }
+    myFixture.checkResultByFile(result);
+  }
+
+  public void testParameter() throws Throwable {
+    doTest("outEmpty/parameter.before.py", "bar", "outEmpty/parameter.after.py");
+  }
+
+  public void testBreakAst() throws Throwable {
+    doTest("outEmpty/break_ast.before.py", "bar", "outEmpty/break_ast.after.py");
+  }
+
+  // TODO[oleg] fix me!!!
+  //public void testExpression() throws Throwable {
+  //  doTest("outEmpty/expression.before.py", "plus", "outEmpty/expression.after.py");
+  //}
+
+  public void testStatement() throws Throwable {
+    doTest("outEmpty/statement.before.py", "foo", "outEmpty/statement.after.py");
+  }
+
+  public void testStatements() throws Throwable {
+    doTest("outEmpty/statements.before.py", "foo", "outEmpty/statements.after.py");
+  }
+
+  // TODO[oleg] fix me!!!
+  //public void testStatementReturn() throws Throwable {
+  //  doTest("outEmpty/statement_return.before.py", "foo", "outEmpty/statement_return.after.py");
+  //}
+
+  public void testBinaryExpression() throws Throwable {
+    doTest("controlFlow/binary_expr.before.py", "foo", "controlFlow/binary_expr.after.py");
+  }
+
+  public void testWhileOutput() throws Throwable {
+    doTest("controlFlow/while_output.before.py", "bar", "controlFlow/while_output.after.py");
+  }
+
+  // TODO[oleg] fix me!!!
+  //public void testComplicated() throws Throwable {
+  //  doTest("controlFlow/complicated.before.py", "foo", "controlFlow/complicated.after.py");
+  //}
+
+  // TODO[oleg] implement me!!!
+  //public void testNameCollisionClass() throws Throwable {
+  //  doTest("namesCollision/class.before.py", "hello", "Method name clashes with already existing method name");
+  //}
+  //
+  //public void testNameCollisionFile() throws Throwable {
+  //  doTest("namesCollision/file.before.py", "hello", "Method name clashes with already existing method name");
+  //}
+  //
+  //public void testNameCollisionSuperClass() throws Throwable {
+  //  doTest("namesCollision/superclass.before.py", "hello", "Method name clashes with already existing method name");
+  //}
+
+  public void testOutNotEmptyStatements() throws Throwable {
+    doTest("outNotEmpty/statements.before.py", "sum_squares", "outNotEmpty/statements.after.py");
+  }
+
+  public void testOutNotEmptyStatements2() throws Throwable {
+    doTest("outNotEmpty/statements2.before.py", "sum_squares", "outNotEmpty/statements2.after.py");
+  }
+
+  public void testComment() throws Throwable {
+    doTest("comment.before.py", "bar", "comment.after.py");
+  }
+}
+