@@ -10,20 +10,20 @@ namespace ts.codefix {
1010 Diagnostics . Argument_of_type_0_is_not_assignable_to_parameter_of_type_1 . code
1111 ] ;
1212
13- enum FixKind {
13+ enum ProblemKind {
1414 MissingReturnStatement ,
1515 MissingParentheses
1616 }
1717
1818 interface MissingReturnInfo {
19- kind : FixKind . MissingReturnStatement ;
19+ kind : ProblemKind . MissingReturnStatement ;
2020 declaration : FunctionLikeDeclaration ;
2121 expression : Expression ;
2222 statement : Statement ;
2323 }
2424
2525 interface MissingParenInfo {
26- kind : FixKind . MissingParentheses ;
26+ kind : ProblemKind . MissingParentheses ;
2727 declaration : ArrowFunction ;
2828 expression : Expression ;
2929 statement : Statement ;
@@ -39,7 +39,7 @@ namespace ts.codefix {
3939 const info = getInfo ( program . getTypeChecker ( ) , sourceFile , start , errorCode ) ;
4040 if ( ! info ) return undefined ;
4141
42- if ( info . kind === FixKind . MissingReturnStatement ) {
42+ if ( info . kind === ProblemKind . MissingReturnStatement ) {
4343 return append (
4444 [ getActionForfixAddReturnStatement ( context , info . expression , info . statement ) ] ,
4545 isArrowFunction ( info . declaration ) ? getActionForfixRemoveBlockBodyBrace ( context , info . declaration , info . expression ) : undefined ) ;
@@ -70,33 +70,14 @@ namespace ts.codefix {
7070 } ) ,
7171 } ) ;
7272
73- function updateFunctionLikeBody ( declaration : FunctionLikeDeclaration , body : Block ) : FunctionLikeDeclaration {
74- switch ( declaration . kind ) {
75- case SyntaxKind . FunctionDeclaration :
76- return createFunctionDeclaration ( declaration . decorators , declaration . modifiers , declaration . asteriskToken , declaration . name , declaration . typeParameters , declaration . parameters , declaration . type , body ) ;
77- case SyntaxKind . MethodDeclaration :
78- return createMethod ( declaration . decorators , declaration . modifiers , declaration . asteriskToken , declaration . name , declaration . questionToken , declaration . typeParameters , declaration . parameters , declaration . type , body ) ;
79- case SyntaxKind . GetAccessor :
80- return createGetAccessor ( declaration . decorators , declaration . modifiers , declaration . name , declaration . parameters , declaration . type , body ) ;
81- case SyntaxKind . SetAccessor :
82- return createSetAccessor ( declaration . decorators , declaration . modifiers , declaration . name , declaration . parameters , body ) ;
83- case SyntaxKind . Constructor :
84- return createConstructor ( declaration . decorators , declaration . modifiers , declaration . parameters , body ) ;
85- case SyntaxKind . FunctionExpression :
86- return createFunctionExpression ( declaration . modifiers , declaration . asteriskToken , declaration . name , declaration . typeParameters , declaration . parameters , declaration . type , body ) ;
87- case SyntaxKind . ArrowFunction :
88- return createArrowFunction ( declaration . modifiers , declaration . typeParameters , declaration . parameters , declaration . type , declaration . equalsGreaterThanToken , body ) ;
89- }
90- }
91-
9273 function getFixInfo ( checker : TypeChecker , declaration : FunctionLikeDeclaration , expectType : Type , isFunctionType : boolean ) : Info | undefined {
9374 if ( ! declaration . body || ! isBlock ( declaration . body ) || length ( declaration . body . statements ) !== 1 ) return undefined ;
9475
9576 const firstStatement = first ( declaration . body . statements ) ;
9677 if ( isExpressionStatement ( firstStatement ) && checkFixedAssignableTo ( checker , declaration , firstStatement . expression , expectType , isFunctionType ) ) {
9778 return {
9879 declaration,
99- kind : FixKind . MissingReturnStatement ,
80+ kind : ProblemKind . MissingReturnStatement ,
10081 expression : firstStatement . expression ,
10182 statement : firstStatement
10283 } ;
@@ -106,12 +87,12 @@ namespace ts.codefix {
10687 if ( checkFixedAssignableTo ( checker , declaration , node , expectType , isFunctionType ) ) {
10788 return isArrowFunction ( declaration ) ? {
10889 declaration,
109- kind : FixKind . MissingParentheses ,
90+ kind : ProblemKind . MissingParentheses ,
11091 expression : node ,
11192 statement : firstStatement
11293 } : {
11394 declaration,
114- kind : FixKind . MissingReturnStatement ,
95+ kind : ProblemKind . MissingReturnStatement ,
11596 expression : node ,
11697 statement : firstStatement
11798 } ;
@@ -124,7 +105,7 @@ namespace ts.codefix {
124105 if ( checkFixedAssignableTo ( checker , declaration , node , expectType , isFunctionType ) ) {
125106 return {
126107 declaration,
127- kind : FixKind . MissingReturnStatement ,
108+ kind : ProblemKind . MissingReturnStatement ,
128109 expression : node ,
129110 statement : firstStatement
130111 } ;
@@ -183,11 +164,16 @@ namespace ts.codefix {
183164 }
184165
185166 function addReturnStatement ( changes : textChanges . ChangeTracker , sourceFile : SourceFile , expression : Expression , statement : Statement ) {
167+ suppressLeadingAndTrailingTrivia ( expression ) ;
186168 changes . replaceNode ( sourceFile , statement , createReturn ( expression ) ) ;
187169 }
188170
189171 function removeBlockBodyBrace ( changes : textChanges . ChangeTracker , sourceFile : SourceFile , declaration : ArrowFunction , expression : Expression , withParen : boolean ) {
190- changes . replaceNode ( sourceFile , declaration . body , ( withParen || needsParentheses ( expression ) ) ? createParen ( expression ) : expression ) ;
172+ const newBody = ( withParen || needsParentheses ( expression ) ) ? createParen ( expression ) : expression ;
173+ suppressLeadingAndTrailingTrivia ( expression ) ;
174+ copyComments ( expression , newBody ) ;
175+
176+ changes . replaceNode ( sourceFile , declaration . body , newBody ) ;
191177 }
192178
193179 function wrapBlockWithParen ( changes : textChanges . ChangeTracker , sourceFile : SourceFile , declaration : ArrowFunction , expression : Expression ) {
@@ -206,6 +192,6 @@ namespace ts.codefix {
206192
207193 function getActionForfixWrapTheBlockWithParen ( context : CodeFixContext , declaration : ArrowFunction , expression : Expression ) {
208194 const changes = textChanges . ChangeTracker . with ( context , t => wrapBlockWithParen ( t , context . sourceFile , declaration , expression ) ) ;
209- return createCodeFixAction ( fixId , changes , Diagnostics . Wrap_this_object_literal_with_parentheses , fixIdWrapTheBlockWithParen , Diagnostics . Wrap_all_object_literal_with_parentheses ) ;
195+ 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 ) ;
210196 }
211197}
0 commit comments