From 9b49498808adf21df9dd4b76e6b11915514e6e72 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Thu, 21 Apr 2016 16:12:20 +0300
Subject: [PATCH] PY-19279 Smart enter on function inserts colon after return
 type annotation

---
 .../smartEnter/fixers/PyFunctionFixer.java    | 13 ++++++++--
 .../colonAfterReturnTypeAnnotation.py         |  1 +
 .../colonAfterReturnTypeAnnotation_after.py   |  2 ++
 .../jetbrains/python/PySmartEnterTest.java    | 26 +++++++------------
 4 files changed, 24 insertions(+), 18 deletions(-)
 create mode 100644 python/testData/codeInsight/smartEnter/colonAfterReturnTypeAnnotation.py
 create mode 100644 python/testData/codeInsight/smartEnter/colonAfterReturnTypeAnnotation_after.py

diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyFunctionFixer.java b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyFunctionFixer.java
index 4bfead4c551b..feec8c681f07 100644
--- a/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyFunctionFixer.java
+++ b/python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyFunctionFixer.java
@@ -20,6 +20,7 @@ import com.intellij.psi.PsiElement;
 import com.intellij.util.IncorrectOperationException;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.codeInsight.editorActions.smartEnter.PySmartEnterProcessor;
+import com.jetbrains.python.psi.PyAnnotation;
 import com.jetbrains.python.psi.PyFunction;
 import com.jetbrains.python.psi.PyParameterList;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
@@ -45,7 +46,15 @@ public class PyFunctionFixer extends PyFixer<PyFunction> {
       if (function.getNameNode() == null) {
         processor.registerUnresolvedError(parameterList.getTextOffset());
       }
-      editor.getDocument().insertString(parameterList.getTextRange().getEndOffset(), ":");
+      final int colonOffset;
+      final PyAnnotation annotation = function.getAnnotation();
+      if (annotation != null) {
+        colonOffset = annotation.getTextRange().getEndOffset();
+      }
+      else {
+        colonOffset = parameterList.getTextRange().getEndOffset();
+      }
+      editor.getDocument().insertString(colonOffset, ":");
     }
   }
 
@@ -54,7 +63,7 @@ public class PyFunctionFixer extends PyFixer<PyFunction> {
    * Attempting to operate in the context of such "fake" function definition may lead to various kinds of malformed code and we want to
    * avoid it.
    *
-   * @return whether it's more the less proper function definition, i.e. it contains at least {@code def} keyword
+   * @return whether it's more or less proper function definition, i.e. it contains at least {@code def} keyword
    */
   static boolean isFakeFunction(@NotNull PyFunction function) {
     return function.getNode().findChildByType(PyTokenTypes.DEF_KEYWORD) == null;
diff --git a/python/testData/codeInsight/smartEnter/colonAfterReturnTypeAnnotation.py b/python/testData/codeInsight/smartEnter/colonAfterReturnTypeAnnotation.py
new file mode 100644
index 000000000000..c74bdadb36cf
--- /dev/null
+++ b/python/testData/codeInsight/smartEnter/colonAfterReturnTypeAnnotation.py
@@ -0,0 +1 @@
+def f() -> int<caret>
\ No newline at end of file
diff --git a/python/testData/codeInsight/smartEnter/colonAfterReturnTypeAnnotation_after.py b/python/testData/codeInsight/smartEnter/colonAfterReturnTypeAnnotation_after.py
new file mode 100644
index 000000000000..aa64571c6301
--- /dev/null
+++ b/python/testData/codeInsight/smartEnter/colonAfterReturnTypeAnnotation_after.py
@@ -0,0 +1,2 @@
+def f() -> int:
+    <caret>
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PySmartEnterTest.java b/python/testSrc/com/jetbrains/python/PySmartEnterTest.java
index f561c1a35a13..761db67c4d49 100644
--- a/python/testSrc/com/jetbrains/python/PySmartEnterTest.java
+++ b/python/testSrc/com/jetbrains/python/PySmartEnterTest.java
@@ -26,6 +26,7 @@ import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
 import com.jetbrains.python.documentation.PyDocumentationSettings;
 import com.jetbrains.python.documentation.docstrings.DocStringFormat;
 import com.jetbrains.python.fixtures.PyTestCase;
+import com.jetbrains.python.psi.LanguageLevel;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.List;
@@ -34,7 +35,7 @@ import java.util.List;
  * @author Alexey.Ivanov
  */
 public class PySmartEnterTest extends PyTestCase {
-  protected static List<SmartEnterProcessor> getSmartProcessors(Language language) {
+  private static List<SmartEnterProcessor> getSmartProcessors(Language language) {
     return SmartEnterProcessors.INSTANCE.forKey(language);
   }
 
@@ -219,29 +220,22 @@ public class PySmartEnterTest extends PyTestCase {
 
   // PY-16765
   public void testGoogleDocStringColonAndIndentAfterSection() {
-    runWithDocStringFormat(DocStringFormat.GOOGLE, new Runnable() {
-      public void run() {
-        doTest();
-      }
-    });
+    runWithDocStringFormat(DocStringFormat.GOOGLE, this::doTest);
   }
 
   // PY-16765
   public void testGoogleDocStringIndentAfterSection() {
-    runWithDocStringFormat(DocStringFormat.GOOGLE, new Runnable() {
-      public void run() {
-        doTest();
-      }
-    });
+    runWithDocStringFormat(DocStringFormat.GOOGLE, this::doTest);
   }
 
   // PY-16765
   public void testGoogleDocStringIndentAfterSectionCustomIndent() {
     getIndentOptions().INDENT_SIZE = 2;
-    runWithDocStringFormat(DocStringFormat.GOOGLE, new Runnable() {
-      public void run() {
-        doTest();
-      }
-    });
+    runWithDocStringFormat(DocStringFormat.GOOGLE, this::doTest);
+  }
+
+  // PY-19279
+  public void testColonAfterReturnTypeAnnotation() {
+    runWithLanguageLevel(LanguageLevel.PYTHON30, this::doTest);
   }
 }