PY-19279 Smart enter on function inserts colon after return type annotation

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;

python/testData/codeInsight/smartEnter/colonAfterReturnTypeAnnotation.py

@@ -0,0 +1 @@
+def f() -> int<caret>

python/testData/codeInsight/smartEnter/colonAfterReturnTypeAnnotation_after.py

@@ -0,0 +1,2 @@
+def f() -> int:
+    <caret>

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);
   }
 }