add support for add or remove braces to arrow function

Author: Kingwl

src/compiler/diagnosticMessages.json

@@ -4165,5 +4165,17 @@
     "Generate 'get' and 'set' accessors": {
         "category": "Message",
         "code": 95046
+    },
+    "Convert arrow function": {
+        "category": "Message",
+        "code": 95047
+    },
+    "Add braces to arrow function": {
+        "category": "Message",
+        "code": 95048
+    },
+    "Remove braces from arrow function": {
+        "category": "Message",
+        "code": 95049
     }
 }

src/harness/tsconfig.json

@@ -114,6 +114,7 @@
         "../services/codefixes/useDefaultImport.ts",
         "../services/refactors/extractSymbol.ts",
         "../services/refactors/generateGetAccessorAndSetAccessor.ts",
+        "../services/refactors/convertArrowFunction.ts",
         "../services/sourcemaps.ts",
         "../services/services.ts",
         "../services/breakpoints.ts",

src/server/tsconfig.json

@@ -110,6 +110,7 @@
         "../services/codefixes/useDefaultImport.ts",
         "../services/refactors/extractSymbol.ts",
         "../services/refactors/generateGetAccessorAndSetAccessor.ts",
+        "../services/refactors/convertArrowFunction.ts",
         "../services/sourcemaps.ts",
         "../services/services.ts",
         "../services/breakpoints.ts",

src/server/tsconfig.library.json

@@ -116,6 +116,7 @@
         "../services/codefixes/useDefaultImport.ts",
         "../services/refactors/extractSymbol.ts",
         "../services/refactors/generateGetAccessorAndSetAccessor.ts",
+        "../services/refactors/convertArrowFunction.ts",
         "../services/sourcemaps.ts",
         "../services/services.ts",
         "../services/breakpoints.ts",

src/services/refactors/convertArrowFunction.ts

@@ -0,0 +1,90 @@
+/* @internal */
+namespace ts.refactor.convertArrowFunction {
+    const refactorName = "Convert arrow function";
+    const refactorDescription = Diagnostics.Convert_arrow_function.message;
+    const addBracesActionName = "Add braces to arrow function";
+    const removeBracesActionName = "Remove braces from arrow function";
+    const addBracesActionDescription = Diagnostics.Add_braces_to_arrow_function.message;
+    const removeBracesActionDescription = Diagnostics.Remove_braces_from_arrow_function.message;
+    registerRefactor(refactorName, { getEditsForAction, getAvailableActions });
+
+    interface Info {
+        container: ArrowFunction;
+        expression: Expression;
+        addBraces: boolean;
+    }
+
+    function getAvailableActions(context: RefactorContext): ApplicableRefactorInfo[] | undefined {
+        const { file, startPosition } = context;
+        const info = getConvertibleArrowFunctionAtPosition(file, startPosition);
+        if (!info) return undefined;
+
+        const actions: RefactorActionInfo[] = [
+            info.addBraces ?
+                {
+                    name: addBracesActionName,
+                    description: addBracesActionDescription
+                } : {
+                    name: removeBracesActionName,
+                    description: removeBracesActionDescription
+                }
+        ];
+
+        return [{
+            name: refactorName,
+            description: refactorDescription,
+            actions
+        }];
+    }
+
+    function getEditsForAction(context: RefactorContext, _actionName: string): RefactorEditInfo | undefined {
+        const { file, startPosition } = context;
+        const info = getConvertibleArrowFunctionAtPosition(file, startPosition);
+        if (!info) return undefined;
+
+        const { addBraces, expression, container } = info;
+        const changeTracker = textChanges.ChangeTracker.fromContext(context);
+        updateBraces(changeTracker, file, container, expression, addBraces);
+
+        return {
+            renameFilename: undefined,
+            renameLocation: undefined,
+            edits: changeTracker.getChanges()
+        };
+    }
+
+    function updateBraces(changeTracker: textChanges.ChangeTracker, file: SourceFile, container: ArrowFunction, expression: Expression, addBraces: boolean) {
+        const body = addBraces ? createBlock([createReturn(expression)]) : expression;
+
+        const arrowFunction = updateArrowFunction(
+            container,
+            container.modifiers,
+            container.typeParameters,
+            container.parameters,
+            container.type,
+            body);
+        changeTracker.replaceNode(file, container, arrowFunction);
+    }
+
+    function getConvertibleArrowFunctionAtPosition(file: SourceFile, startPosition: number): Info | undefined {
+        const node = getTokenAtPosition(file, startPosition, /*includeJsDocComment*/ false);
+        const container = getContainingFunction(node);
+        if (!container || !isArrowFunction(container)) return undefined;
+
+        if (isExpression(container.body)) {
+            return {
+                container,
+                addBraces: true,
+                expression: container.body
+            };
+        }
+        else if (container.body.statements.length === 1 && isReturnStatement(first(container.body.statements))) {
+            return {
+                container,
+                addBraces: false,
+                expression: (<ReturnStatement>first(container.body.statements)).expression
+            };
+        }
+        return undefined;
+    }
+}

src/services/tsconfig.json

@@ -107,6 +107,7 @@
         "codefixes/useDefaultImport.ts",
         "refactors/extractSymbol.ts",
         "refactors/generateGetAccessorAndSetAccessor.ts",
+        "refactors/convertArrowFunction.ts",
         "sourcemaps.ts",
         "services.ts",
         "breakpoints.ts",

tests/cases/fourslash/refactorAddBracesToArrowFunction1.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => a + 1;
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Convert arrow function",
+    actionName: "Add braces to arrow function",
+    actionDescription: "Add braces to arrow function",
+    newContent: `const foo = a => { return a + 1; };`,
+});

tests/cases/fourslash/refactorAddBracesToArrowFunction2.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => ({ a: 1 });
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Convert arrow function",
+    actionName: "Add braces to arrow function",
+    actionDescription: "Add braces to arrow function",
+    newContent: `const foo = a => { return ({ a: 1 }); };`,
+});

tests/cases/fourslash/refactorAddBracesToArrowFunction3.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => 1;
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Convert arrow function",
+    actionName: "Add braces to arrow function",
+    actionDescription: "Add braces to arrow function",
+    newContent: `const foo = a => { return 1; };`,
+});

tests/cases/fourslash/refactorAddBracesToArrowFunction4.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return a + 1; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Convert arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => a + 1;`,
+});

tests/cases/fourslash/refactorAddBracesToArrowFunction5.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return { a: 1 }; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Convert arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => ({ a: 1 });`,
+});

tests/cases/fourslash/refactorAddBracesToArrowFunction6.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return 1; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Convert arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => 1;`,
+});

tests/cases/fourslash/refactorAddBracesToArrowFunction7.ts

@@ -0,0 +1,6 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => {  };
+
+goTo.select("a", "b");
+verify.not.refactorAvailable("Convert arrow function");

tests/cases/fourslash/refactorAddBracesToArrowFunction8.ts

@@ -0,0 +1,9 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => {
+////     const b = 1;
+////     return a + b;
+//// };
+
+goTo.select("a", "b");
+verify.not.refactorAvailable("Convert arrow function");