@@ -17,8 +17,7 @@ namespace ts.refactor.inferFunctionReturnType {
1717 function getEditsForAction ( context : RefactorContext ) : RefactorEditInfo | undefined {
1818 const info = getInfo ( context ) ;
1919 if ( info && ! isRefactorErrorInfo ( info ) ) {
20- const edits = textChanges . ChangeTracker . with ( context , t =>
21- t . tryInsertTypeAnnotation ( context . file , info . declaration , info . returnTypeNode ) ) ;
20+ const edits = textChanges . ChangeTracker . with ( context , t => doChange ( context . file , t , info . declaration , info . returnTypeNode ) ) ;
2221 return { renameFilename : undefined , renameLocation : undefined , edits } ;
2322 }
2423 return undefined ;
@@ -55,6 +54,19 @@ namespace ts.refactor.inferFunctionReturnType {
5554 returnTypeNode : TypeNode ;
5655 }
5756
57+ function doChange ( sourceFile : SourceFile , changes : textChanges . ChangeTracker , declaration : ConvertibleDeclaration , typeNode : TypeNode ) {
58+ const closeParen = findChildOfKind ( declaration , SyntaxKind . CloseParenToken , sourceFile ) ;
59+ const needParens = isArrowFunction ( declaration ) && closeParen === undefined ;
60+ const endNode = needParens ? first ( declaration . parameters ) : closeParen ;
61+ if ( endNode ) {
62+ if ( needParens ) {
63+ changes . insertNodeBefore ( sourceFile , endNode , factory . createToken ( SyntaxKind . OpenParenToken ) ) ;
64+ changes . insertNodeAfter ( sourceFile , endNode , factory . createToken ( SyntaxKind . CloseParenToken ) ) ;
65+ }
66+ changes . insertNodeAt ( sourceFile , endNode . end , typeNode , { prefix : ": " } ) ;
67+ }
68+ }
69+
5870 function getInfo ( context : RefactorContext ) : FunctionInfo | RefactorErrorInfo | undefined {
5971 if ( isInJSFile ( context . file ) || ! refactorKindBeginsWith ( inferReturnTypeAction . kind , context . kind ) ) return ;
6072
0 commit comments