Repository: nicoespeon/abracadabra
Pull request #748 (2022-11-05 14:39:08)

Change signature III: Handle class method
Another one.

*This PR contains commits from* :  https://github.com/nicoespeon/abracadabra/pull/747

There are numerations for PR's Orders (I, II, III) _sorry combining commigs from differents PRS..._

Playground example

![class-method](https://user-images.githubusercontent.com/8685132/200125173-1e95c493-1bee-4b6f-a265-dc4f2805138c.gif)


For Objects propertys I will search a good solution for it in a other iteration (I will take a break now :P) 




Patch:
diff --git a/src/playground/maths/calculator.ts b/src/playground/maths/calculator.ts
index 19fb18e9..81491a0e 100644
--- a/src/playground/maths/calculator.ts
+++ b/src/playground/maths/calculator.ts
@@ -1,6 +1,10 @@
-import { add } from "./maths";
+import { add, subtract, Maths } from "./maths";
 
 add(3, 400);
+subtract(5, 500);
+
+const math = new Maths();
+math.add(0, 1);
 
 // 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..8fb25019 100644
--- a/src/playground/maths/maths.ts
+++ b/src/playground/maths/maths.ts
@@ -2,4 +2,15 @@ export function add(param: number, num2: number) {
   return param + num2;
 }
 
+export const subtract = (param: number, num2: number) => {
+  return param - num2;
+};
+
+export class Maths {
+  add(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..ca6af08e 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,10 @@ function updateCode(
           return t.identifier("undefined");
         });
         node.arguments = newArgs;
-      } else if (t.isFunctionDeclaration(node)) {
+      } else if (
+        isFunctionDeclarationOrArrowFunction(node) ||
+        t.isClassMethod(node)
+      ) {
         const params = node.params.slice();
         if (params.length) {
           newPositions.forEach((order) => {
@@ -135,13 +147,31 @@ function updateCode(
 
 export function createVisitor(
   selection: Selection,
-  onMatch: (path: t.NodePath<t.FunctionDeclaration>) => void
+  onMatch: (
+    path: t.NodePath<
+      t.FunctionDeclaration | t.ArrowFunctionExpression | t.ClassMethod
+    >,
+    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));
+    },
+    ClassMethod(path) {
+      if (!selection.isInsidePath(path)) return;
+      onMatch(path, selection);
     }
   };
 }
@@ -209,6 +239,15 @@ function createVisitorForReferences(
     FunctionDeclaration(path) {
       if (!selection.isInsidePath(path)) return;
       onMatch(path);
+    },
+    ArrowFunctionExpression(path) {
+      if (!selection.isInsidePath(path)) return;
+
+      onMatch(path);
+    },
+    ClassMethod(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..2bd7fac1 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,118 @@ 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);
+    }
+  );
+
+  testEach<{
+    code: Code;
+    expected: Code;
+  }>(
+    "In same file with class methods",
+    [
+      {
+        description: "when there are a class without references",
+        code: `class Maths {
+        [cursor]add(a, b) {
+          return a + b;
+        }
+      }`,
+        expected: `class Maths {
+        add(b, a) {
+          return a + b;
+        }
+      }`
+      },
+      {
+        description: "when there are a class with method references",
+        code: `class Maths {
+        [cursor]add(a, b) {
+          return a + b;
+        }
+      }
+      const maths = new Maths();
+      maths.add(1, 2);`,
+        expected: `class Maths {
+        add(b, a) {
+          return a + b;
+        }
+      }
+      const maths = new Maths();
+      maths.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 +363,116 @@ 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
+              }
+            ]
+          }
+        },
+        {
+          description: "that import a class and use a method",
+          setup: {
+            currentFile: {
+              code: `export class Maths {
+                [cursor]add(a, b) {
+                  return a + b;
+                }
+              }`,
+              path: new AbsolutePath("/temp/module.ts")
+            },
+            otherFiles: [
+              {
+                code: `import {Maths} from './module';
+                const maths = new Maths();
+                maths.add(1, 2);
+              `,
+                path: addModule
+              },
+              {
+                code: `import {add} from './anotherModule';
+                export const calculateAdd = (a, b) => {
+                  const maths = new Maths();
+                  return maths.add(a, b);
+                }
+              `,
+                path: anotherModule
+              }
+            ]
+          },
+          expected: {
+            currentFile: {
+              code: `export class Maths {
+                add(b, a) {
+                  return a + b;
+                }
+              }`,
+              path: new AbsolutePath("/temp/module.ts")
+            },
+            otherFiles: [
+              {
+                code: `import {Maths} from './module';
+                const maths = new Maths();
+                maths.add(2, 1);
+              `,
+                path: addModule
+              },
+              {
+                code: `import {add} from './anotherModule';
+                export const calculateAdd = (a, b) => {
+                  const maths = new Maths();
+                  return maths.add(b, a);
+                }
+              `,
+                path: anotherModule
+              }
+            ]
+          }
         }
       ],
       async ({ setup, expected }) => {