refactor and handle comment

microsoft · Mar 26, 2020 · 16d90f0 · 16d90f0
1 parent af9552e
commit 16d90f0
Show file tree

Hide file tree

Showing 11 changed files with 84 additions and 55 deletions.
diff --git a/src/compiler/diagnosticMessages.json b/src/compiler/diagnosticMessages.json
@@ -5561,7 +5561,7 @@
         "category": "Message",
         "code": 95109
     },
-    "Wrap this object literal with parentheses": {
+    "Wrap the following body with parentheses which should be an object literal": {
         "category": "Message",
         "code": 95110
     },

diff --git a/src/compiler/factoryPublic.ts b/src/compiler/factoryPublic.ts
@@ -2042,6 +2042,25 @@ namespace ts {
             : node;
     }
 
+    export function updateFunctionLikeBody(declaration: FunctionLikeDeclaration, body: Block): FunctionLikeDeclaration {
+        switch (declaration.kind) {
+            case SyntaxKind.FunctionDeclaration:
+                return createFunctionDeclaration(declaration.decorators, declaration.modifiers, declaration.asteriskToken, declaration.name, declaration.typeParameters, declaration.parameters, declaration.type, body);
+            case SyntaxKind.MethodDeclaration:
+                return createMethod(declaration.decorators, declaration.modifiers, declaration.asteriskToken, declaration.name, declaration.questionToken, declaration.typeParameters, declaration.parameters, declaration.type, body);
+            case SyntaxKind.GetAccessor:
+                return createGetAccessor(declaration.decorators, declaration.modifiers, declaration.name, declaration.parameters, declaration.type, body);
+            case SyntaxKind.SetAccessor:
+                return createSetAccessor(declaration.decorators, declaration.modifiers, declaration.name, declaration.parameters, body);
+            case SyntaxKind.Constructor:
+                return createConstructor(declaration.decorators, declaration.modifiers, declaration.parameters, body);
+            case SyntaxKind.FunctionExpression:
+                return createFunctionExpression(declaration.modifiers, declaration.asteriskToken, declaration.name, declaration.typeParameters, declaration.parameters, declaration.type, body);
+            case SyntaxKind.ArrowFunction:
+                return createArrowFunction(declaration.modifiers, declaration.typeParameters, declaration.parameters, declaration.type, declaration.equalsGreaterThanToken, body);
+        }
+    }
+
     export function createClassDeclaration(
         decorators: readonly Decorator[] | undefined,
         modifiers: readonly Modifier[] | undefined,

diff --git a/src/services/codefixes/returnValueCorrect.ts b/src/services/codefixes/returnValueCorrect.ts
@@ -10,20 +10,20 @@ namespace ts.codefix {
         Diagnostics.Argument_of_type_0_is_not_assignable_to_parameter_of_type_1.code
     ];
 
-    enum FixKind {
+    enum ProblemKind {
         MissingReturnStatement,
         MissingParentheses
     }
 
     interface MissingReturnInfo {
-        kind: FixKind.MissingReturnStatement;
+        kind: ProblemKind.MissingReturnStatement;
         declaration: FunctionLikeDeclaration;
         expression: Expression;
         statement: Statement;
     }
 
     interface MissingParenInfo {
-        kind: FixKind.MissingParentheses;
+        kind: ProblemKind.MissingParentheses;
         declaration: ArrowFunction;
         expression: Expression;
         statement: Statement;
@@ -39,7 +39,7 @@ namespace ts.codefix {
             const info = getInfo(program.getTypeChecker(), sourceFile, start, errorCode);
             if (!info) return undefined;
 
-            if (info.kind === FixKind.MissingReturnStatement) {
+            if (info.kind === ProblemKind.MissingReturnStatement) {
                 return append(
                     [getActionForfixAddReturnStatement(context, info.expression, info.statement)],
                     isArrowFunction(info.declaration) ? getActionForfixRemoveBlockBodyBrace(context, info.declaration, info.expression): undefined);
@@ -70,33 +70,14 @@ namespace ts.codefix {
         }),
     });
 
-    function updateFunctionLikeBody(declaration: FunctionLikeDeclaration, body: Block): FunctionLikeDeclaration {
-        switch (declaration.kind) {
-            case SyntaxKind.FunctionDeclaration:
-                return createFunctionDeclaration(declaration.decorators, declaration.modifiers, declaration.asteriskToken, declaration.name, declaration.typeParameters, declaration.parameters, declaration.type, body);
-            case SyntaxKind.MethodDeclaration:
-                return createMethod(declaration.decorators, declaration.modifiers, declaration.asteriskToken, declaration.name, declaration.questionToken, declaration.typeParameters, declaration.parameters, declaration.type, body);
-            case SyntaxKind.GetAccessor:
-                return createGetAccessor(declaration.decorators, declaration.modifiers, declaration.name, declaration.parameters, declaration.type, body);
-            case SyntaxKind.SetAccessor:
-                return createSetAccessor(declaration.decorators, declaration.modifiers, declaration.name, declaration.parameters, body);
-            case SyntaxKind.Constructor:
-                return createConstructor(declaration.decorators, declaration.modifiers, declaration.parameters, body);
-            case SyntaxKind.FunctionExpression:
-                return createFunctionExpression(declaration.modifiers, declaration.asteriskToken, declaration.name, declaration.typeParameters, declaration.parameters, declaration.type, body);
-            case SyntaxKind.ArrowFunction:
-                return createArrowFunction(declaration.modifiers, declaration.typeParameters, declaration.parameters, declaration.type, declaration.equalsGreaterThanToken, body);
-        }
-    }
-
     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) && checkFixedAssignableTo(checker, declaration, firstStatement.expression, expectType, isFunctionType)) {
             return {
                 declaration,
-                kind: FixKind.MissingReturnStatement,
+                kind: ProblemKind.MissingReturnStatement,
                 expression: firstStatement.expression,
                 statement: firstStatement
             };
@@ -106,12 +87,12 @@ namespace ts.codefix {
             if (checkFixedAssignableTo(checker, declaration, node, expectType, isFunctionType)) {
                 return isArrowFunction(declaration) ? {
                     declaration,
-                    kind: FixKind.MissingParentheses,
+                    kind: ProblemKind.MissingParentheses,
                     expression: node,
                     statement: firstStatement
                 } : {
                         declaration,
-                        kind: FixKind.MissingReturnStatement,
+                        kind: ProblemKind.MissingReturnStatement,
                         expression: node,
                         statement: firstStatement
                     };
@@ -124,7 +105,7 @@ namespace ts.codefix {
                 if (checkFixedAssignableTo(checker, declaration, node, expectType, isFunctionType)) {
                     return {
                         declaration,
-                        kind: FixKind.MissingReturnStatement,
+                        kind: ProblemKind.MissingReturnStatement,
                         expression: node,
                         statement: firstStatement
                     };
@@ -183,11 +164,16 @@ namespace ts.codefix {
     }
 
     function addReturnStatement(changes: textChanges.ChangeTracker, sourceFile: SourceFile, expression: Expression, statement: Statement) {
+        suppressLeadingAndTrailingTrivia(expression);
         changes.replaceNode(sourceFile, statement, createReturn(expression));
     }
 
     function removeBlockBodyBrace(changes: textChanges.ChangeTracker, sourceFile: SourceFile, declaration: ArrowFunction, expression: Expression, withParen: boolean) {
-        changes.replaceNode(sourceFile, declaration.body, (withParen || needsParentheses(expression)) ? createParen(expression) : expression);
+        const newBody = (withParen || needsParentheses(expression)) ? createParen(expression) : expression;
+        suppressLeadingAndTrailingTrivia(expression);
+        copyComments(expression, newBody);
+
+        changes.replaceNode(sourceFile, declaration.body, newBody);
     }
 
     function wrapBlockWithParen(changes: textChanges.ChangeTracker, sourceFile: SourceFile, declaration: ArrowFunction, expression: Expression) {
@@ -206,6 +192,6 @@ namespace ts.codefix {
 
     function getActionForfixWrapTheBlockWithParen(context: CodeFixContext, declaration: ArrowFunction, expression: Expression) {
         const changes = textChanges.ChangeTracker.with(context, t => wrapBlockWithParen(t, context.sourceFile, declaration, expression));
-        return createCodeFixAction(fixId, changes, Diagnostics.Wrap_this_object_literal_with_parentheses, fixIdWrapTheBlockWithParen, Diagnostics.Wrap_all_object_literal_with_parentheses);
+        return createCodeFixAction(fixId, changes, Diagnostics.Wrap_the_following_body_with_parentheses_which_should_be_an_object_literal, fixIdWrapTheBlockWithParen, Diagnostics.Wrap_all_object_literal_with_parentheses);
     }
 }
diff --git a/src/services/refactors/convertParamsToDestructuredObject.ts b/src/services/refactors/convertParamsToDestructuredObject.ts
@@ -491,27 +491,6 @@ namespace ts.refactor.convertParamsToDestructuredObject {
         }
     }
 
-    function copyComments(sourceNode: Node, targetNode: Node) {
-        const sourceFile = sourceNode.getSourceFile();
-        const text = sourceFile.text;
-        if (hasLeadingLineBreak(sourceNode, text)) {
-            copyLeadingComments(sourceNode, targetNode, sourceFile);
-        }
-        else {
-            copyTrailingAsLeadingComments(sourceNode, targetNode, sourceFile);
-        }
-        copyTrailingComments(sourceNode, targetNode, sourceFile);
-    }
-
-    function hasLeadingLineBreak(node: Node, text: string) {
-        const start = node.getFullStart();
-        const end = node.getStart();
-        for (let i = start; i < end; i++) {
-            if (text.charCodeAt(i) === CharacterCodes.lineFeed) return true;
-        }
-        return false;
-    }
-
     function getParameterName(paramDeclaration: ValidParameterDeclaration) {
         return getTextOfIdentifierOrLiteral(paramDeclaration.name);
     }

diff --git a/src/services/utilities.ts b/src/services/utilities.ts
@@ -2250,6 +2250,27 @@ namespace ts {
         addEmitFlagsRecursively(node, EmitFlags.NoTrailingComments, getLastChild);
     }
 
+    export function copyComments(sourceNode: Node, targetNode: Node) {
+        const sourceFile = sourceNode.getSourceFile();
+        const text = sourceFile.text;
+        if (hasLeadingLineBreak(sourceNode, text)) {
+            copyLeadingComments(sourceNode, targetNode, sourceFile);
+        }
+        else {
+            copyTrailingAsLeadingComments(sourceNode, targetNode, sourceFile);
+        }
+        copyTrailingComments(sourceNode, targetNode, sourceFile);
+    }
+
+    function hasLeadingLineBreak(node: Node, text: string) {
+        const start = node.getFullStart();
+        const end = node.getStart();
+        for (let i = start; i < end; i++) {
+            if (text.charCodeAt(i) === CharacterCodes.lineFeed) return true;
+        }
+        return false;
+    }
+
     function addEmitFlagsRecursively(node: Node, flag: EmitFlags, getChild: (n: Node) => Node | undefined) {
         addEmitFlags(node, flag);
         const child = getChild(node);

diff --git a/tests/cases/fourslash/codeFixCorrectReturnValue14.ts b/tests/cases/fourslash/codeFixCorrectReturnValue14.ts
@@ -8,6 +8,6 @@
 //// }
 
 verify.codeFixAvailable([
-    { description: 'Wrap this object literal with parentheses' },
+    { description: 'Wrap the following body with parentheses which should be an object literal' },
     { description: 'Remove unused label' },
 ]);
diff --git a/tests/cases/fourslash/codeFixCorrectReturnValue15.ts b/tests/cases/fourslash/codeFixCorrectReturnValue15.ts
@@ -7,6 +7,6 @@
 //// Foo(() => { bar: '1' })
 
 verify.codeFixAvailable([
-    { description: 'Wrap this object literal with parentheses' },
+    { description: 'Wrap the following body with parentheses which should be an object literal' },
     { description: 'Remove unused label' },
 ]);
diff --git a/tests/cases/fourslash/codeFixCorrectReturnValue16.ts b/tests/cases/fourslash/codeFixCorrectReturnValue16.ts
@@ -8,6 +8,6 @@
 //// Foo(() => { bar: '1' })
 
 verify.codeFixAvailable([
-    { description: 'Wrap this object literal with parentheses' },
+    { description: 'Wrap the following body with parentheses which should be an object literal' },
     { description: 'Remove unused label' },
 ]);
diff --git a/tests/cases/fourslash/codeFixCorrectReturnValue19.ts b/tests/cases/fourslash/codeFixCorrectReturnValue19.ts
@@ -14,7 +14,7 @@
 //// var x = <Comp t={() => { bar: '1' }} />;
 
 verify.codeFixAvailable([
-    { description: 'Wrap this object literal with parentheses' },
+    { description: 'Wrap the following body with parentheses which should be an object literal' },
     { description: `Infer type of 'props' from usage` },
     { description: 'Remove unused label' },
 ]);
diff --git a/tests/cases/fourslash/codeFixCorrectReturnValue24.ts b/tests/cases/fourslash/codeFixCorrectReturnValue24.ts
@@ -0,0 +1,12 @@
+/// <reference path='fourslash.ts' />
+
+//// function Foo (a: () => number) { a() }
+//// Foo(() => { /* leading */ 1 /* trailing */ })
+
+verify.codeFix({
+    description: "Add a return statement",
+    index: 0,
+    newFileContent: 
+    `function Foo (a: () => number) { a() }
+Foo(() => { /* leading */ return 1 /* trailing */ })`
+})
diff --git a/tests/cases/fourslash/codeFixCorrectReturnValue25.ts b/tests/cases/fourslash/codeFixCorrectReturnValue25.ts
@@ -0,0 +1,12 @@
+/// <reference path='fourslash.ts' />
+
+//// function Foo (a: () => number) { a() }
+//// Foo(() => { /* leading */ 1 /* trailing */ })
+
+verify.codeFix({
+    description: "Remove block body braces",
+    index: 1,
+    newFileContent: 
+    `function Foo (a: () => number) { a() }
+Foo(() => /* leading */ 1 /* trailing */)`
+})