add support for more case

microsoft · amcasey · Apr 2, 2020 · Aug 10, 2018 · Aug 14, 2018 · Aug 14, 2018
commit 0b1e6cc6ea15df4a599d317775bfa9a52885fb78
diff --git a/src/services/codefixes/returnValueSurmise.ts b/src/services/codefixes/returnValueSurmise.ts
@@ -16,11 +16,18 @@ namespace ts.codefix {
         MissingParentheses
     }
 
-    interface Info {
+    interface MissingReturnInfo {
+        kind: FixKind.MissingReturnStatement;
         declaration: FunctionLikeDeclaration;
-        fixKind: FixKind;
     }
 
+    interface MissingParenthesesInfo {
+        kind: FixKind.MissingParentheses;
+        declaration: ArrowFunction;
+    }
+
+    type Info = MissingReturnInfo | MissingParenthesesInfo;
+
     registerCodeFix({
         errorCodes,
         fixIds: [fixIdAddReturnStatement, fixIdRemoveBlockBodyBrace, fixIdReplaceBraceWithParen, fixIdWrapTheBlockWithParen],
@@ -29,7 +36,7 @@ namespace ts.codefix {
             const info = getInfo(program.getTypeChecker(), sourceFile, start);
             if (!info) return undefined;
 
-            if (info.fixKind === FixKind.MissingReturnStatement) {
+            if (info.kind === FixKind.MissingReturnStatement) {
                 return [
                     getActionForfixAddReturnStatement(context, info.declaration),
                     getActionForfixRemoveBlockBodyBrace(context, info.declaration),
@@ -55,7 +62,7 @@ namespace ts.codefix {
                     removeBlockBodyBrace(changes, diag.file, info.declaration, /* withParen */ true);
                     break;
                 case fixIdWrapTheBlockWithParen:
-                    wrapBlockWithParen(changes, diag.file, info.declaration);
+                    wrapBlockWithParen(changes, diag.file, <ArrowFunction>info.declaration);
                     break;
                 default:
                     Debug.fail(JSON.stringify(context.fixId));
@@ -82,61 +89,43 @@ namespace ts.codefix {
         }
     }
 
-    function getFixInfo(checker: TypeChecker, declaration: FunctionLikeDeclaration, expectType: Type): Info | undefined {
+    function getFixInfo(checker: TypeChecker, declaration: FunctionLikeDeclaration, expectType: Type, isFunctionType: boolean): Info | undefined {
         if (!declaration.body || !isBlock(declaration.body) || length(declaration.body.statements) !== 1) return undefined;
 
         const firstStatement = first(declaration.body.statements);
-        if (isExpressionStatement(firstStatement)) {
-            const maybeFunction = updateFunctionLikeBody(declaration, createBlock([createReturn(firstStatement.expression)]));
-            const maybeType = checker.getTypeAtLocation(maybeFunction);
-            if (checker.isTypeAssignableTo(maybeType, expectType)) {
+        if (isExpressionStatement(firstStatement) && checkFixedAssignableTo(checker, declaration, firstStatement.expression, expectType, isFunctionType)) {
+            return {
+                declaration,
+                kind: FixKind.MissingReturnStatement
+            };
+        }
+        else if (isArrowFunction(declaration) && isLabeledStatement(firstStatement) && isExpressionStatement(firstStatement.statement)) {
+            const node = createObjectLiteral([createPropertyAssignment(firstStatement.label, firstStatement.statement.expression)]);
+            if (checkFixedAssignableTo(checker, declaration, node, expectType, isFunctionType)) {
                 return {
                     declaration,
-                    fixKind: FixKind.MissingReturnStatement
+                    kind: FixKind.MissingParentheses
                 };
             }
         }
-        else if (isLabeledStatement(firstStatement) && isExpressionStatement(firstStatement.statement)) {
-            const maybeFunction = updateFunctionLikeBody(declaration, createBlock([createReturn(
-                createObjectLiteral([createPropertyAssignment(firstStatement.label, firstStatement.statement.expression)])
-            )]));
-            const maybeType = checker.getTypeAtLocation(maybeFunction);
-            if (checker.isTypeAssignableTo(maybeType, expectType)) {
-                return {
-                    declaration,
-                    fixKind: FixKind.MissingParentheses
-                };
+        else if (isBlock(firstStatement) && length(firstStatement.statements) === 1) {
+            const firstBlockStatement = first(firstStatement.statements);
+            if (isLabeledStatement(firstBlockStatement) && isExpressionStatement(firstBlockStatement.statement)) {
+                const node = createObjectLiteral([createPropertyAssignment(firstBlockStatement.label, firstBlockStatement.statement.expression)]);
+                if (checkFixedAssignableTo(checker, declaration, node, expectType, isFunctionType)) {
+                    return {
+                        declaration,
+                        kind: FixKind.MissingReturnStatement
+                    };
+                }
             }
         }
 
         return undefined;
     }
 
-    function getFixInfoFromTypeAnnotation(checker: TypeChecker, declaration: FunctionLikeDeclaration, expectType: Type): Info | undefined {
-        if (!declaration.body || !isBlock(declaration.body) || length(declaration.body.statements) !== 1) return undefined;
-
-        const firstStatement = first(declaration.body.statements);
-        if (isExpressionStatement(firstStatement)) {
-            const actualType = checker.getTypeAtLocation(firstStatement.expression);
-            if (checker.isTypeAssignableTo(actualType, expectType)) {
-                return {
-                    declaration,
-                    fixKind: FixKind.MissingReturnStatement
-                };
-            }
-        }
-        else if (isLabeledStatement(firstStatement) && isExpressionStatement(firstStatement.statement)) {
-            const node = createObjectLiteral([createPropertyAssignment(firstStatement.label, firstStatement.statement.expression)]);
-            const type = checker.getTypeAtLocation(node);
-            if (checker.isTypeAssignableTo(type, expectType)) {
-                return {
-                    declaration,
-                    fixKind: FixKind.MissingParentheses
-                };
-            }
-        }
-
-        return undefined;
+    function checkFixedAssignableTo (checker: TypeChecker, declaration: FunctionLikeDeclaration, expr: Expression, type: Type, isFunctionType: boolean) {
+        return checker.isTypeAssignableTo(checker.getTypeAtLocation(isFunctionType ? updateFunctionLikeBody(declaration, createBlock([createReturn(expr)])) : expr), type);
     }
 
     function getInfo(checker: TypeChecker, sourceFile: SourceFile, position: number): Info | undefined {
@@ -146,41 +135,54 @@ namespace ts.codefix {
         const declaration = findAncestor(node.parent, isFunctionLikeDeclaration);
         if (declaration) {
             if (declaration.body && declaration.type && rangeContainsRange(declaration.type, node)) {
-                // check by type annotation
-                return getFixInfoFromTypeAnnotation(checker, declaration, checker.getTypeFromTypeNode(declaration.type));
+                return getFixInfo(checker, declaration, checker.getTypeFromTypeNode(declaration.type), /* isFunctionType */ false);
             }
             else if (isCallExpression(declaration.parent) && declaration.body) {
                 const pos = declaration.parent.arguments.indexOf(<Expression>declaration);
                 const type = checker.getContextualTypeForArgumentAtIndex(declaration.parent, pos);
                 if (!type) return undefined;
-                return getFixInfo(checker, declaration, type);
+                return getFixInfo(checker, declaration, type, /* isFunctionType */ true);
             }
         }
         else if (isDeclarationName(node) && isVariableDeclaration(node.parent) && node.parent.initializer && node.parent.type) {
             const initializer = skipParentheses(node.parent.initializer);
             if (!isFunctionLikeDeclaration(initializer) || !initializer.body) return undefined;
-            return getFixInfo(checker, initializer, checker.getTypeFromTypeNode(node.parent.type));
+            return getFixInfo(checker, initializer, checker.getTypeFromTypeNode(node.parent.type), /* isFunctionType */ true);
+        }
+        return Debug.fail("unknow pattern");
+    }
+
+    function getReturnExpression (stmt: Statement) {
+        if (isExpressionStatement(stmt)) {
+            return stmt.expression;
         }
-        Debug.fail("unknow pattern");
+        else if (isBlock(stmt) && length(stmt.statements) === 1) {
+            const block = first(stmt.statements);
+            if (isLabeledStatement(block) && isExpressionStatement(block.statement)) {
+                return block.statement.expression;
+            }
+        }
+
+        return Debug.fail("unknow statement");
     }
 
     function addReturnStatement(changes: textChanges.ChangeTracker, sourceFile: SourceFile, declaration: FunctionLikeDeclaration) {
         const body = cast(declaration.body, isBlock);
-        const firstStatement = cast(first(body.statements), isExpressionStatement);
-        changes.replaceNode(sourceFile, firstStatement, createReturn(firstStatement.expression));
+        const firstStatement = first(body.statements);
+        changes.replaceNode(sourceFile, firstStatement, createReturn(getReturnExpression(firstStatement)));
     }
 
     function removeBlockBodyBrace(changes: textChanges.ChangeTracker, sourceFile: SourceFile, declaration: FunctionLikeDeclaration, withParen: boolean) {
         const body = cast(declaration.body, isBlock);
-        const expression = cast(first(body.statements), isExpressionStatement).expression;
+        const expression = getReturnExpression(first(body.statements));
         changes.replaceNode(sourceFile, body, withParen ? createParen(expression) : expression);
     }
 
-    function wrapBlockWithParen(changes: textChanges.ChangeTracker, sourceFile: SourceFile, declaration: FunctionLikeDeclaration) {
+    function wrapBlockWithParen(changes: textChanges.ChangeTracker, sourceFile: SourceFile, declaration: ArrowFunction) {
         const body = cast(declaration.body, isBlock);
         const labeledStatement = cast(first(body.statements), isLabeledStatement);
         const expression = cast(labeledStatement.statement, isExpressionStatement).expression;
-        changes.replaceNode(sourceFile, body, createParen(createObjectLiteral([createPropertyAssignment(labeledStatement.label, expression)])));
+        changes.replaceNode(sourceFile, declaration.body, createParen(createObjectLiteral([createPropertyAssignment(labeledStatement.label, expression)])));
     }
 
     function getActionForfixAddReturnStatement(context: CodeFixContext, declaration: FunctionLikeDeclaration) {
@@ -198,7 +200,7 @@ namespace ts.codefix {
         return createCodeFixAction(fixId, changes, Diagnostics.Replace_braces_with_parentheses, fixIdReplaceBraceWithParen, Diagnostics.Surmise_all_return_value);
     }
 
-    function getActionForfixWrapTheBlockWithParen(context: CodeFixContext, declaration: FunctionLikeDeclaration) {
+    function getActionForfixWrapTheBlockWithParen(context: CodeFixContext, declaration: ArrowFunction) {
         const changes = textChanges.ChangeTracker.with(context, t => wrapBlockWithParen(t, context.sourceFile, declaration));
         return createCodeFixAction(fixId, changes, Diagnostics.Wrap_this_block_with_parentheses, fixIdWrapTheBlockWithParen, Diagnostics.Surmise_all_return_value);
     }

diff --git a/tests/cases/fourslash/codeFixSurmiseReturnValue11.ts b/tests/cases/fourslash/codeFixSurmiseReturnValue11.ts
@@ -0,0 +1,15 @@
+/// <reference path='fourslash.ts' />
+//// interface A {
+////     bar: string
+//// }
+//// 
+//// function Foo (): A {
+////     { bar: '123' }
+//// }
+
+verify.codeFixAvailable([
+    { description: 'Add a return statement' },
+    { description: 'Remove block body braces' },
+    { description: 'Replace braces with parentheses' },
+    { description: 'Remove unused label' },
+]);
diff --git a/tests/cases/fourslash/codeFixSurmiseReturnValue12.ts b/tests/cases/fourslash/codeFixSurmiseReturnValue12.ts
@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+//// interface A {
+////     bar: string
+//// }
+//// 
+//// function Foo (a: () => A) { a() }
+//// Foo(() => {  
+////     { bar: '123' }
+//// })
+
+verify.codeFixAvailable([
+    { description: 'Add a return statement' },
+    { description: 'Remove block body braces' },
+    { description: 'Replace braces with parentheses' },
+    { description: 'Remove unused label' },
+]);
diff --git a/tests/cases/fourslash/codeFixSurmiseReturnValue13.ts b/tests/cases/fourslash/codeFixSurmiseReturnValue13.ts
@@ -0,0 +1,15 @@
+/// <reference path='fourslash.ts' />
+//// interface A {
+////     bar: string
+//// }
+//// 
+//// const a: () => A = () => {
+////     { bar: '1' }    
+//// }
+
+verify.codeFixAvailable([
+    { description: 'Add a return statement' },
+    { description: 'Remove block body braces' },
+    { description: 'Replace braces with parentheses' },
+    { description: 'Remove unused label' },
+]);
diff --git a/tests/cases/fourslash/codeFixSurmiseReturnValue14.ts b/tests/cases/fourslash/codeFixSurmiseReturnValue14.ts
@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+//// interface A {
+////     bar: string
+//// }
+//// 
+//// const a: () => A = () => {
+////     bar: '1'   
+//// }
+
+verify.codeFixAvailable([
+    { description: 'Wrap this block with parentheses' },
+    { description: 'Remove unused label' },
+]);
diff --git a/tests/cases/fourslash/codeFixSurmiseReturnValue15.ts b/tests/cases/fourslash/codeFixSurmiseReturnValue15.ts
@@ -0,0 +1,12 @@
+/// <reference path='fourslash.ts' />
+//// interface A {
+////     bar: string
+//// }
+//// 
+//// function Foo (a: () => A) { a() }
+//// Foo(() => { bar: '1' })
+
+verify.codeFixAvailable([
+    { description: 'Wrap this block with parentheses' },
+    { description: 'Remove unused label' },
+]);