feat(40750): add refactoring to infer a return type annotation to a function

Author: a-tarasyuk

src/compiler/diagnosticMessages.json

@@ -5955,6 +5955,10 @@
         "category": "Message",
         "code": 95144
     },
+    "Infer function return type": {
+        "category": "Message",
+        "code": 95145
+    },
 
     "No value exists in scope for the shorthand property '{0}'. Either declare one or provide an initializer.": {
         "category": "Error",

src/services/refactors/inferFunctionReturnType.ts

@@ -0,0 +1,72 @@
+/* @internal */
+namespace ts.refactor.inferFunctionReturnType {
+    const refactorName = "Infer function return type";
+    const refactorDescription = Diagnostics.Infer_function_return_type.message;
+    registerRefactor(refactorName, { getEditsForAction, getAvailableActions });
+
+    function getEditsForAction(context: RefactorContext): RefactorEditInfo | undefined {
+        const info = getInfo(context);
+        if (info) {
+            const edits = textChanges.ChangeTracker.with(context, t =>
+                t.tryInsertTypeAnnotation(context.file, info.declaration, info.returnType));
+            return { renameFilename: undefined, renameLocation: undefined, edits };
+        }
+        return undefined;
+    }
+
+    function getAvailableActions(context: RefactorContext): readonly ApplicableRefactorInfo[] {
+        const info = getInfo(context);
+        if (info) {
+            return [{
+                name: refactorName,
+                description: refactorDescription,
+                actions: [{
+                    name: refactorName,
+                    description: refactorDescription
+                }]
+            }];
+        }
+        return emptyArray;
+    }
+
+    type ConvertableDeclaration =
+        | FunctionDeclaration
+        | FunctionExpression
+        | ArrowFunction
+        | MethodDeclaration;
+
+    interface Info {
+        declaration: ConvertableDeclaration;
+        returnType: TypeNode;
+    }
+
+    function getInfo(context: RefactorContext): Info | undefined {
+        if (isInJSFile(context.file)) return;
+
+        const token = getTokenAtPosition(context.file, context.startPosition);
+        const declaration = findAncestor(token, isConvertableDeclaration);
+        if (!declaration || declaration.type) return;
+
+        const typeChecker = context.program.getTypeChecker();
+        const type = typeChecker.getTypeAtLocation(declaration);
+        const signatures = type.getCallSignatures();
+        if (!signatures.length) return;
+
+        const returnType = typeChecker.typeToTypeNode(first(signatures).getReturnType(), declaration, NodeBuilderFlags.NoTruncation);
+        if (returnType) {
+            return { declaration, returnType };
+        }
+    }
+
+    function isConvertableDeclaration(node: Node): node is ConvertableDeclaration {
+        switch (node.kind) {
+            case SyntaxKind.FunctionDeclaration:
+            case SyntaxKind.FunctionExpression:
+            case SyntaxKind.ArrowFunction:
+            case SyntaxKind.MethodDeclaration:
+                return true;
+            default:
+                return false;
+        }
+    }
+}

src/services/tsconfig.json

@@ -120,6 +120,7 @@
         "refactors/convertParamsToDestructuredObject.ts",
         "refactors/convertStringOrTemplateLiteral.ts",
         "refactors/convertArrowFunctionOrFunctionExpression.ts",
+        "refactors/inferFunctionReturnType.ts",
         "services.ts",
         "breakpoints.ts",
         "transform.ts",

tests/cases/fourslash/refactorInferFunctionReturnType1.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////function /*a*/foo/*b*/() {
+////    return { x: 1, y: 1 };
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`function foo(): { x: number; y: number; } {
+    return { x: 1, y: 1 };
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType10.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////function /*a*/foo/*b*/(x: number) {
+////    return x ? x : x > 1;
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`function foo(x: number): number | boolean {
+    return x ? x : x > 1;
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType11.ts

@@ -0,0 +1,36 @@
+/// <reference path='fourslash.ts' />
+
+////interface F1 { x: number; y: number; }
+////type T1 = [number, number];
+////
+////function /*a*/foo/*b*/(num: number) {
+////   switch (num) {
+////      case 1:
+////         return { x: num, y: num } as F1;
+////      case 2:
+////         return [num, num] as T1;
+////      default:
+////         return num;
+////   }
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`interface F1 { x: number; y: number; }
+type T1 = [number, number];
+
+function foo(num: number): number | F1 | T1 {
+   switch (num) {
+      case 1:
+         return { x: num, y: num } as F1;
+      case 2:
+         return [num, num] as T1;
+      default:
+         return num;
+   }
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType2.ts

@@ -0,0 +1,20 @@
+/// <reference path='fourslash.ts' />
+
+////class Foo {
+////    /*a*/method/*b*/() {
+////        return { x: 1, y: 1 };
+////    }
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`class Foo {
+    method(): { x: number; y: number; } {
+        return { x: 1, y: 1 };
+    }
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType3.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////const foo = /*a*/function/*b*/() {
+////    return { x: 1, y: 1 };
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`const foo = function(): { x: number; y: number; } {
+    return { x: 1, y: 1 };
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType4.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////const foo = /*a*/()/*b*/ => {
+////    return { x: 1, y: 1 };
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`const foo = (): { x: number; y: number; } => {
+    return { x: 1, y: 1 };
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType5.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////function /*a*/foo/*b*/() {
+////    return 1;
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`function foo(): number {
+    return 1;
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType6.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////function /*a*/foo/*b*/() {
+////    return "";
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`function foo(): string {
+    return "";
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType7.ts

@@ -0,0 +1,14 @@
+/// <reference path='fourslash.ts' />
+
+////function /*a*/foo/*b*/() {
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`function foo(): void {
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType8.ts

@@ -0,0 +1,18 @@
+/// <reference path='fourslash.ts' />
+
+////function /*a*/foo/*b*/() {
+////    const bar = 1 as any;
+////    return bar;
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`function foo(): any {
+    const bar = 1 as any;
+    return bar;
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType9.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////function /*a*/foo<T>/*b*/() {
+////    return 1 as T;
+////}
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`function foo<T>(): T {
+    return 1 as T;
+}`
+});