PY-17094 Quickfix that adds Py3 return type annotation handles compound statements...

diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java
index c6f13a67839f4d3a7c8e993f9d625c4af723fd85..dd3f1aae4f75cc40cbf415914fae326a5f17bf31 100644 (file)
--- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java
@@ -17,19 +17,23 @@ package com.jetbrains.python.codeInsight.intentions;
 
 import com.intellij.codeInsight.CodeInsightUtilCore;
 import com.intellij.codeInsight.template.*;
+import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.fileEditor.FileEditorManager;
 import com.intellij.openapi.fileEditor.OpenFileDescriptor;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.psi.PsiDocumentManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiReference;
 import com.intellij.util.IncorrectOperationException;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.PyNames;
+import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyPsiUtils;
 import org.jetbrains.annotations.NotNull;
 
 /**
@@ -68,7 +72,7 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
       annotateParameter(project, editor, parameter);
     }
     else {
-      annotateReturnType(project, elementAt);
+      annotateReturnType(project, editor.getDocument(), elementAt);
     }
   }
 
@@ -99,24 +103,32 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
     }
   }
 
-  private void annotateReturnType(Project project, PsiElement resolved) {
-    final PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
-
+  private void annotateReturnType(Project project, Document document, PsiElement resolved) {
     PyCallable callable = getCallable(resolved);
 
     if (callable instanceof PyFunction) {
-      final String functionSignature = "def " + callable.getName() + callable.getParameterList().getText();
-      String functionText = functionSignature +
-                            " -> object:";
-      final PyStatementList statementList = ((PyFunction)callable).getStatementList();
-      for (PyStatement st : statementList.getStatements()) {
-        functionText = functionText + "\n\t" + st.getText();
+      final String annotationText = " -> " + PyNames.OBJECT;
+      
+      final PsiElement prevElem = PyPsiUtils.getPrevNonCommentSibling(((PyFunction)callable).getStatementList(), true);
+      assert prevElem != null;
+
+      final PsiDocumentManager manager = PsiDocumentManager.getInstance(project);
+      try {
+        final TextRange range = prevElem.getTextRange();
+        manager.doPostponedOperationsAndUnblockDocument(document);
+        if (prevElem.getNode().getElementType() == PyTokenTypes.COLON) {
+          document.insertString(range.getStartOffset(), annotationText);
+        }
+        else {
+          document.insertString(range.getEndOffset(), annotationText + ":");
+        }
       }
-      final PyFunction function = elementGenerator.createFromText(LanguageLevel.forElement(callable), PyFunction.class,
-                                                                  functionText);
-      callable = (PyFunction)callable.replace(function);
+      finally {
+        manager.commitDocument(document);
+      }
+      
+      
       callable = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(callable);
-
       final PyAnnotation annotation = ((PyFunction)callable).getAnnotation();
       assert annotation != null;
       final PyExpression annotationValue = annotation.getValue();

diff --git a/python/testData/intentions/afterReturnTypeInPy3AnnotationLocalFunction.py b/python/testData/intentions/afterReturnTypeInPy3AnnotationLocalFunction.py
new file mode 100644 (file)
index 0000000..9d43307
--- /dev/null
+++ b/python/testData/intentions/afterReturnTypeInPy3AnnotationLocalFunction.py
@@ -0,0 +1,4 @@
+def foo() -> object:
+    @decorator
+    def bar():
+        pass
\ No newline at end of file

diff --git a/python/testData/intentions/afterReturnTypeInPy3AnnotationNoColon.py b/python/testData/intentions/afterReturnTypeInPy3AnnotationNoColon.py
new file mode 100644 (file)
index 0000000..9d43307
--- /dev/null
+++ b/python/testData/intentions/afterReturnTypeInPy3AnnotationNoColon.py
@@ -0,0 +1,4 @@
+def foo() -> object:
+    @decorator
+    def bar():
+        pass
\ No newline at end of file

diff --git a/python/testData/intentions/beforeReturnTypeInPy3AnnotationLocalFunction.py b/python/testData/intentions/beforeReturnTypeInPy3AnnotationLocalFunction.py
new file mode 100644 (file)
index 0000000..e3c6d23
--- /dev/null
+++ b/python/testData/intentions/beforeReturnTypeInPy3AnnotationLocalFunction.py
@@ -0,0 +1,4 @@
+def fo<caret>o():
+    @decorator
+    def bar():
+        pass
\ No newline at end of file

diff --git a/python/testData/intentions/beforeReturnTypeInPy3AnnotationNoColon.py b/python/testData/intentions/beforeReturnTypeInPy3AnnotationNoColon.py
new file mode 100644 (file)
index 0000000..2ea3c54
--- /dev/null
+++ b/python/testData/intentions/beforeReturnTypeInPy3AnnotationNoColon.py
@@ -0,0 +1,4 @@
+def fo<caret>o()
+    @decorator
+    def bar():
+        pass
\ No newline at end of file

diff --git a/python/testSrc/com/jetbrains/python/intentions/PyIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyIntentionTest.java
index 7d23e79313ffbaea85535404bc198894ec7418b1..3d218448a462ebf47749684262b032242ae28ca2 100644 (file)
--- a/python/testSrc/com/jetbrains/python/intentions/PyIntentionTest.java
+++ b/python/testSrc/com/jetbrains/python/intentions/PyIntentionTest.java
@@ -359,6 +359,15 @@ public class  PyIntentionTest extends PyTestCase {
   public void testReturnTypeInPy3Annotation2() {      //PY-8783
     doTest(PyBundle.message("INTN.specify.return.type.in.annotation"), LanguageLevel.PYTHON32);
   }
+  
+  // PY-17094
+  public void testReturnTypeInPy3AnnotationLocalFunction() {
+    doTest(PyBundle.message("INTN.specify.return.type.in.annotation"), LanguageLevel.PYTHON32);
+  }
+  
+  public void testReturnTypeInPy3AnnotationNoColon() {
+    doTest(PyBundle.message("INTN.specify.return.type.in.annotation"), LanguageLevel.PYTHON32);
+  }
 
   public void testTypeAnnotation3() {  //PY-7087
     doTypeAnnotationTest();