Repository: nicoespeon/abracadabra
Pull request #747 (2022-11-05 13:56:11)

Change signature II: Handle arrow functions 
Hey :) 

**This PR depends on https://github.com/nicoespeon/abracadabra/pull/746**

I was able to change function signature of an arrow function, but there are some doubts when to trigger "Change signature" action.

In this iteration we shoud put the Cursor on function asignation:

![Untitled-2022-02-17-2107](https://user-images.githubusercontent.com/8685132/200123081-f9e5d3d7-0c62-4e5c-b016-fc856136df4f.png)

But maybe we want to trigger "Change signature" over variable declaration. 

![here-trigger](https://user-images.githubusercontent.com/8685132/200123208-b33a5a45-509e-4c16-a1af-1be17625d870.png)

**Where do you think we should trigger it** @nicoespeon ?


nicoespeon: @11joselu I think it's fine to restrict the scope of the trigger to the function params itself for now 👍 
We'll see from usage if it feels like it would be intuitive to have it on the variable assignment too. But I think it's fair to put your cursor on the params if that's what you want to change.
nicoespeon: @11joselu could you update this PR with the latest main? It includes part 1 + a little change in the contract test that causes a conflict.

Patch:
diff --git a/src/playground/maths/calculator.ts b/src/playground/maths/calculator.ts
index 19fb18e9..df4b67e5 100644
--- a/src/playground/maths/calculator.ts
+++ b/src/playground/maths/calculator.ts
@@ -1,6 +1,7 @@
-import { add } from "./maths";
+import { add, subtract } from "./maths";
 
 add(3, 400);
+subtract(5, 500);
 
 // Used to view what params generates in webview table
 function viewParametersTable(
diff --git a/src/playground/maths/maths.ts b/src/playground/maths/maths.ts
index eb2d00b8..878650f4 100644
--- a/src/playground/maths/maths.ts
+++ b/src/playground/maths/maths.ts
@@ -2,4 +2,9 @@ export function add(param: number, num2: number) {
   return param + num2;
 }
 
+export const subtract = (param: number, num2: number) => {
+  return param - num2;
+};
+
 add(1, 200);
+subtract(3, 300);
diff --git a/src/refactorings/change-signature/change-signature.ts b/src/refactorings/change-signature/change-signature.ts
index e8ce6f42..f87cf3c6 100644
--- a/src/refactorings/change-signature/change-signature.ts
+++ b/src/refactorings/change-signature/change-signature.ts
@@ -2,10 +2,11 @@ import * as t from "../../ast";
 import { Editor, ErrorReason, SelectedPosition } from "../../editor/editor";
 import { Path } from "../../editor/path";
 import { Selection } from "../../editor/selection";
+import { isFunctionDeclarationOrArrowFunction } from "../../ast/identity";
 
 export async function changeSignature(editor: Editor) {
   const { code, selection } = editor;
-  const params = getParams(code, selection);
+  const { fixedSelection, params } = getParams(code, selection);
 
   if (!params) {
     editor.showError(ErrorReason.CantChangeSignature);
@@ -13,8 +14,7 @@ export async function changeSignature(editor: Editor) {
   }
 
   await editor.askForPositions(params, async (newPositions) => {
-    const { selection } = editor;
-    const refrences = await editor.getSelectionReferences(selection);
+    const refrences = await editor.getSelectionReferences(fixedSelection);
 
     const filesContent = await Promise.all(
       refrences.map(async (reference) => {
@@ -69,12 +69,16 @@ export async function changeSignature(editor: Editor) {
 
 type Params = { label: string; value: { startAt: number; endAt: number } }[];
 
-function getParams(code: string, selection: Selection): Params | null {
+function getParams(
+  code: string,
+  selection: Selection
+): { params: Params | null; fixedSelection: Selection } {
   let result: Params | null = null;
+  let arrowSelection: Selection = selection;
 
   t.parseAndTraverseCode(
     code,
-    createVisitor(selection, (path) => {
+    createVisitor(selection, (path, aArrowSelection) => {
       result = path.node.params.map((p, index) => {
         return {
           label: getParamName(p),
@@ -84,11 +88,16 @@ function getParams(code: string, selection: Selection): Params | null {
           }
         };
       });
+
+      arrowSelection = aArrowSelection;
       path.stop();
     })
   );
 
-  return result;
+  return {
+    params: result,
+    fixedSelection: arrowSelection
+  };
 }
 
 function updateCode(
@@ -116,7 +125,7 @@ function updateCode(
           return t.identifier("undefined");
         });
         node.arguments = newArgs;
-      } else if (t.isFunctionDeclaration(node)) {
+      } else if (isFunctionDeclarationOrArrowFunction(node)) {
         const params = node.params.slice();
         if (params.length) {
           newPositions.forEach((order) => {
@@ -135,13 +144,25 @@ function updateCode(
 
 export function createVisitor(
   selection: Selection,
-  onMatch: (path: t.NodePath<t.FunctionDeclaration>) => void
+  onMatch: (
+    path: t.NodePath<t.FunctionDeclaration | t.ArrowFunctionExpression>,
+    arrowSelection: Selection
+  ) => void
 ): t.Visitor {
   return {
     FunctionDeclaration(path) {
       if (!selection.isInsidePath(path)) return;
 
-      onMatch(path);
+      onMatch(path, selection);
+    },
+    ArrowFunctionExpression(path) {
+      if (!selection.isInsidePath(path)) return;
+
+      if (!t.isVariableDeclarator(path.parent)) return;
+
+      if (!path.parent.loc) return;
+
+      onMatch(path, Selection.fromAST(path.parent.loc));
     }
   };
 }
@@ -208,6 +229,11 @@ function createVisitorForReferences(
     },
     FunctionDeclaration(path) {
       if (!selection.isInsidePath(path)) return;
+      onMatch(path);
+    },
+    ArrowFunctionExpression(path) {
+      if (!selection.isInsidePath(path)) return;
+
       onMatch(path);
     }
   };


Test patch:
diff --git a/src/refactorings/change-signature/change-signature.test.ts b/src/refactorings/change-signature/change-signature.test.ts
index 9b653f3a..f31110f5 100644
--- a/src/refactorings/change-signature/change-signature.test.ts
+++ b/src/refactorings/change-signature/change-signature.test.ts
@@ -14,7 +14,7 @@ describe("Change Signature", () => {
     code: Code;
     expected: Code;
   }>(
-    "In same file",
+    "In the same file with function declarations",
     [
       {
         description: "when there are a function without references",
@@ -185,6 +185,67 @@ describe("Change Signature", () => {
     }
   );
 
+  testEach<{
+    code: Code;
+    expected: Code;
+  }>(
+    "In the same file with arrow function declarations",
+    [
+      {
+        description: "when there is an arrow function without references",
+        code: `const add = [cursor](a, b) => {
+            return a + b;
+          }`,
+        expected: `const add = (b, a) => {
+            return a + b;
+          }`
+      },
+      {
+        description: "when there is an arrow function with references",
+        code: `const add = [cursor](a, b) => {
+            return a + b;
+          }
+          add(1, 2);
+          `,
+        expected: `const add = (b, a) => {
+            return a + b;
+          }
+          add(2, 1);
+          `
+      },
+      {
+        description:
+          "when there is a defined arrow function in multiple lines with references",
+        code: `const add = [cursor](
+          a,
+          b) => {
+            return a + b;
+          }
+          add(1, 2);
+          `,
+        expected: `const add = (
+          b,
+          a) => {
+            return a + b;
+          }
+          add(2, 1);
+          `
+      }
+    ],
+    async ({ code, expected }) => {
+      const path = new AbsolutePath("/temp/aFile.ts");
+      const editor = new InMemoryEditor(code);
+      await editor.writeIn(path, editor.code);
+      editor.saveUserChoices(userChangePositionOf(0, 1));
+      editor.saveUserChoices(userChangePositionOf(1, 0));
+
+      await changeSignature(editor);
+
+      const extracted = await editor.codeOf(path);
+      expect(extracted).toBe(expected);
+    }
+  );
+
   describe("Modules", () => {
     const addModule = new AbsolutePath("/temp/add.ts");
     const anotherModule = new AbsolutePath("/temp/anotherModule");
@@ -251,6 +312,57 @@ describe("Change Signature", () => {
               }
             ]
           }
+        },
+        {
+          description: "that import arrow function",
+          setup: {
+            currentFile: {
+              code: `export const add = [cursor](a, b) => {
+              return a + b;
+            }`,
+              path: new AbsolutePath("/temp/module.ts")
+            },
+            otherFiles: [
+              {
+                code: `import {add} from './module';
+                add(1, 2)
+              `,
+                path: addModule
+              },
+              {
+                code: `import {add} from './anotherModule';
+                export const calculateAdd = (a, b) => {
+                  return add(a, b);
+                }
+              `,
+                path: anotherModule
+              }
+            ]
+          },
+          expected: {
+            currentFile: {
+              code: `export const add = (b, a) => {
+              return a + b;
+            }`,
+              path: new AbsolutePath("/temp/module.ts")
+            },
+            otherFiles: [
+              {
+                code: `import {add} from './module';
+                add(2, 1)
+              `,
+                path: addModule
+              },
+              {
+                code: `import {add} from './anotherModule';
+                export const calculateAdd = (a, b) => {
+                  return add(b, a);
+                }
+              `,
+                path: anotherModule
+              }
+            ]
+          }
         }
       ],
       async ({ setup, expected }) => {