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

src/compiler/diagnosticMessages.json

@@ -5979,6 +5979,10 @@
         "category": "Message",
         "code": 95147
     },
+    "Infer function return type": {
+        "category": "Message",
+        "code": 95148
+    },
 
     "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,84 @@
+/* @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.returnTypeNode));
+            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 ConvertibleDeclaration =
+        | FunctionDeclaration
+        | FunctionExpression
+        | ArrowFunction
+        | MethodDeclaration;
+
+    interface Info {
+        declaration: ConvertibleDeclaration;
+        returnTypeNode: TypeNode;
+    }
+
+    function getInfo(context: RefactorContext): Info | undefined {
+        if (isInJSFile(context.file)) return;
+
+        const token = getTokenAtPosition(context.file, context.startPosition);
+        const declaration = findAncestor(token, isConvertibleDeclaration);
+        if (!declaration || !declaration.body || declaration.type) return;
+
+        const typeChecker = context.program.getTypeChecker();
+        const returnType = tryGetReturnType(typeChecker, declaration);
+        if (!returnType) return;
+
+        const returnTypeNode = typeChecker.typeToTypeNode(returnType, declaration, NodeBuilderFlags.NoTruncation);
+        if (returnTypeNode) {
+            return { declaration, returnTypeNode };
+        }
+    }
+
+    function isConvertibleDeclaration(node: Node): node is ConvertibleDeclaration {
+        switch (node.kind) {
+            case SyntaxKind.FunctionDeclaration:
+            case SyntaxKind.FunctionExpression:
+            case SyntaxKind.ArrowFunction:
+            case SyntaxKind.MethodDeclaration:
+                return true;
+            default:
+                return false;
+        }
+    }
+
+    function tryGetReturnType(typeChecker: TypeChecker, node: ConvertibleDeclaration): Type | undefined {
+        if (typeChecker.isImplementationOfOverload(node)) {
+            const signatures = typeChecker.getTypeAtLocation(node).getCallSignatures();
+            if (signatures.length > 1) {
+                return typeChecker.getUnionType(mapDefined(signatures, s => s.getReturnType()));
+            }
+        }
+        const signature = typeChecker.getSignatureFromDeclaration(node);
+        if (signature) {
+            return typeChecker.getReturnTypeOfSignature(signature);
+        }
+    }
+}

src/services/tsconfig.json

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

tests/cases/fourslash/refactorConvertExport_namedToDefault_alreadyHasDefault.ts

@@ -5,4 +5,4 @@
 ////export default function g() {}
 
 goTo.select("a", "b");
-verify.refactorsAvailable([]);
+verify.refactorsAvailable(["Infer function return type"]);

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/refactorInferFunctionReturnType12.ts

@@ -0,0 +1,22 @@
+/// <reference path='fourslash.ts' />
+
+////function f(x: number) {
+////    return 1;
+////}
+////function /*a*/f/*b*/(x: number) {
+////    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 f(x: number) {
+    return 1;
+}
+function f(x: number): string {
+    return "1";
+}`
+});

tests/cases/fourslash/refactorInferFunctionReturnType13.ts

@@ -0,0 +1,9 @@
+/// <reference path='fourslash.ts' />
+
+////function /*a*/f/*b*/(x: string): number;
+////function f(x: string | number) {
+////  return 1;
+////}
+
+goTo.select("a", "b");
+verify.not.refactorAvailable("Infer function return type");

tests/cases/fourslash/refactorInferFunctionReturnType14.ts

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

tests/cases/fourslash/refactorInferFunctionReturnType15.ts

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

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