correctly handle return annotations in 'override' and formatter (PY-2690)

python/src/com/jetbrains/python/codeInsight/override/PyOverrideImplementUtil.java

@@ -149,6 +149,10 @@ public class PyOverrideImplementUtil {
     if (decorators != null && decorators.findDecorator(PyNames.CLASSMETHOD) != null) {
       pyFunctionBuilder.decorate(PyNames.CLASSMETHOD);
     }
+    PyAnnotation anno = baseFunction.getAnnotation();
+    if (anno != null) {
+      pyFunctionBuilder.annotation(anno.getText());
+    }
     final PyParameter[] baseParams = baseFunction.getParameterList().getParameters();
     for (PyParameter parameter : baseParams) {
       pyFunctionBuilder.parameter(parameter.getText());

python/src/com/jetbrains/python/formatter/PyBlock.java

@@ -277,6 +277,18 @@ public class PyBlock implements ASTBlock {
       return Spacing.createSpacing(0, Integer.MAX_VALUE, 1, false, 1);
     }
 
+    if (parentType == PyElementTypes.ANNOTATION) {
+      if (type1 == PyTokenTypes.GT) {
+        return createSpaces(1);
+      }
+      if (type1 == PyTokenTypes.MINUS && type2 == PyTokenTypes.GT) {
+        return createSpaces(0);
+      }
+    }
+    if (type2 == PyElementTypes.ANNOTATION) {
+      return createSpaces(1);
+    }
+
     if (type1 == PyTokenTypes.COLON) {
       if (type2 == PyElementTypes.STATEMENT_LIST) {
         return Spacing.createSpacing(1, Integer.MAX_VALUE, 0, true, 0);
@@ -338,7 +350,7 @@ public class PyBlock implements ASTBlock {
       if (parentType == PyElementTypes.NAMED_PARAMETER ||
           parentType == PyElementTypes.STAR_ARGUMENT_EXPRESSION ||
           parentType == PyElementTypes.STAR_EXPRESSION) {
-        return Spacing.createSpacing(0, 0, 0, mySettings.KEEP_LINE_BREAKS, mySettings.KEEP_BLANK_LINES_IN_CODE);
+        return createSpaces(0);
       }
       return getSpacingForOption(mySettings.SPACE_AROUND_MULTIPLICATIVE_OPERATORS);
     }
@@ -392,8 +404,11 @@ public class PyBlock implements ASTBlock {
   }
 
   private Spacing getSpacingForOption(boolean isOptionSet) {
-    int spaces = isOptionSet ? 1 : 0;
-    return Spacing.createSpacing(spaces, spaces, 0, mySettings.KEEP_LINE_BREAKS, mySettings.KEEP_BLANK_LINES_IN_CODE);
+    return createSpaces(isOptionSet ? 1 : 0);
+  }
+
+  private Spacing createSpaces(int count) {
+    return Spacing.createSpacing(count, count, 0, mySettings.KEEP_LINE_BREAKS, mySettings.KEEP_BLANK_LINES_IN_CODE);
   }
 
   private static boolean isStatementOrDeclaration(final IElementType type) {

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

@@ -22,6 +22,7 @@ public class PyFunctionBuilder {
   private final List<String> myParameters = new ArrayList<String>();
   private final List<String> myStatements = new ArrayList<String>();
   private final List<String> myDecorators = new ArrayList<String>();
+  private String myAnnotation = null;
 
   public PyFunctionBuilder(String name) {
     myName = name;
@@ -38,6 +39,11 @@ public class PyFunctionBuilder {
     return this;
   }
 
+  public PyFunctionBuilder annotation(String text) {
+    myAnnotation = text;
+    return this;
+  }
+
   public PyFunctionBuilder statement(String text) {
     myStatements.add(text);
     return this;
@@ -65,7 +71,11 @@ public class PyFunctionBuilder {
     builder.append("def ");
     builder.append(myName).append("(");
     builder.append(StringUtil.join(myParameters, ", "));
-    builder.append("):");
+    builder.append(")");
+    if (myAnnotation != null) {
+      builder.append(myAnnotation);
+    }
+    builder.append(":");
     List<String> statements = myStatements.isEmpty() ? Collections.singletonList("pass") : myStatements;
     final CodeStyleSettings codeStyleSettings = CodeStyleSettingsManager.getInstance(project).getCurrentSettings();
     int indentSize = codeStyleSettings.getIndentOptions(PythonFileType.INSTANCE).INDENT_SIZE;

python/testData/override/returnAnnotation.py

@@ -0,0 +1,6 @@
+class A():
+    def some_method(self) -> "return value":
+        pass
+
+class B(A):
+    pass

python/testData/override/returnAnnotation_after.py

@@ -0,0 +1,8 @@
+class A():
+    def some_method(self) -> "return value":
+        pass
+
+class B(A):
+    def some_method(self) -> "return value":
+        <selection>super().some_method()</selection>
+

python/testSrc/com/jetbrains/python/PyOverrideTest.java

@@ -78,8 +78,11 @@ public class PyOverrideTest extends PyLightFixtureTestCase {
     doTest3k();
   }
 
-
   public void testTypeAnnotations() {  // PY-2547
     doTest3k();
   }
+
+  public void testReturnAnnotation() {  // PY-2690
+    doTest3k();
+  }
 }