done

packages/cli/src/helper/file.ts

@@ -16,12 +16,26 @@ export function removeIndexesFromSourceCode(
 ) {
   let newSourceCode = sourceCode;
 
-  // sort to start removing from the of the file end
-  indexesToRemove.sort((a, b) => b.startIndex - a.startIndex);
+  // Sort the indexes based on startIndex (ascending)
+  indexesToRemove.sort((a, b) => a.startIndex - b.startIndex);
+
+  // Merge overlapping or contiguous ranges
+  const mergedIndexes: { startIndex: number; endIndex: number }[] = [];
+  for (const range of indexesToRemove) {
+    const last = mergedIndexes[mergedIndexes.length - 1];
+    if (last && range.startIndex <= last.endIndex) {
+      // Merge the current range with the last range
+      last.endIndex = Math.max(last.endIndex, range.endIndex);
+    } else {
+      // Add as a new range
+      mergedIndexes.push({ ...range });
+    }
+  }
 
-  // TODO improve this method by implementing merging of indexes to prevent accidental deletion of code
+  // Sort to start removing from the end of the file
+  mergedIndexes.sort((a, b) => b.startIndex - a.startIndex);
 
-  indexesToRemove.forEach(({ startIndex, endIndex }) => {
+  mergedIndexes.forEach(({ startIndex, endIndex }) => {
     newSourceCode =
       newSourceCode.slice(0, startIndex) + newSourceCode.slice(endIndex);
   });

packages/cli/src/languagesPlugins/javascript.test.ts

@@ -750,10 +750,10 @@ foofoo();
 foobar();
 `,
       expectedSourceCode: `
-import foo, { a, c as c_n,  } from "./export1";
+import foo, { a, c as c_n, z } from "./export1";
 import barbar from "./export2";
-
-
+import foofoo from "./export3";
+import foobar from "./export4";
 
 foo();
 a();
@@ -781,10 +781,10 @@ let foofoo: FooFoo;
 let foobar: FooBar;
 `,
       expectedSourceCode: `
-import Foo, { A, C as C_N,  } from "./export1";
+import Foo, { A, C as C_N, Z } from "./export1";
 import BarBar from "./export2";
-
-
+import FooFoo from "./export3";
+import FooBar from "./export4";
 
 let foo: Foo;
 let a: A;

packages/cli/src/languagesPlugins/javascript.ts

@@ -59,8 +59,12 @@ class JavascriptPlugin implements LanguagePlugin {
           return;
         }
 
-        // remove the other annotations
-        const nextNode = node.nextNamedSibling;
+        let nextNode = node.nextNamedSibling;
+        // We need to remove all decorators too
+        while (nextNode && nextNode.type === "decorator") {
+          nextNode = nextNode.nextNamedSibling;
+        }
+
         if (!nextNode) {
           throw new Error("Could not find next node");
         }
@@ -475,14 +479,8 @@ class JavascriptPlugin implements LanguagePlugin {
     return depExports;
   }
 
-  // TODO thought, we could leverage LLM to better find the usages of the import identifiers
-  // This is really a case by case issue here. This implementation only works for some cases
-  #getImportIdentifiersUsages(
-    node: Parser.SyntaxNode,
-    identifier: Parser.SyntaxNode,
-  ) {
-    let usageNodes: Parser.SyntaxNode[] = [];
-    const identifierQuery = new Parser.Query(
+  #getIdentifiersNode(node: Parser.SyntaxNode, identifier: Parser.SyntaxNode) {
+    const query = new Parser.Query(
       this.parser.getLanguage(),
       this.isTypescript
         ? `
@@ -499,13 +497,26 @@ class JavascriptPlugin implements LanguagePlugin {
       `,
     );
 
-    const identifierCaptures = identifierQuery.captures(node);
-    identifierCaptures.forEach((capture) => {
-      if (capture.node.id === identifier.id) {
-        return;
-      }
+    const captures = query.captures(node);
+    const nodes = captures.map((capture) => capture.node);
+
+    const identifiers = nodes.filter((node) => node.id !== identifier.id);
+
+    return identifiers;
+  }
+
+  // TODO We could leverage LLM to better find the usages of the import identifiers
+  // This is really a case by case issue here. This implementation only works for some cases
+  #getImportIdentifiersUsages(
+    node: Parser.SyntaxNode,
+    identifier: Parser.SyntaxNode,
+  ) {
+    const otherIdentifiers = this.#getIdentifiersNode(node, identifier);
+
+    const usageNodes: Parser.SyntaxNode[] = [];
 
-      let targetNode = capture.node;
+    otherIdentifiers.forEach((otherIdentifier) => {
+      let targetNode = otherIdentifier;
       while (true) {
         // we can remove from the array
         if (targetNode.parent && targetNode.parent.type === "array") {
@@ -530,10 +541,6 @@ class JavascriptPlugin implements LanguagePlugin {
       return usageNodes.push(targetNode);
     });
 
-    usageNodes = usageNodes.filter(
-      (usageNode) => usageNode.id !== identifier.id,
-    );
-
     return usageNodes;
   }
 
@@ -544,7 +551,6 @@ class JavascriptPlugin implements LanguagePlugin {
   ) {
     const indexesToRemove: { startIndex: number; endIndex: number }[] = [];
 
-    // Parse the source code to get the AST
     const tree = this.parser.parse(sourceCode);
 
     // Get all imports from the file
@@ -565,12 +571,6 @@ class JavascriptPlugin implements LanguagePlugin {
       const depExport = depExportMap.get(depImport.source);
 
       if (!depExport) {
-        // If no exports exist for the import source, mark the entire import for removal
-        indexesToRemove.push({
-          startIndex: depImport.node.startIndex,
-          endIndex: depImport.node.endIndex,
-        });
-
         // Remove usages of all identifiers from this import
         depImport.identifiers.forEach((importIdentifier) => {
           const usageNodes = this.#getImportIdentifiersUsages(
@@ -642,22 +642,6 @@ class JavascriptPlugin implements LanguagePlugin {
           });
         }
       });
-
-      if (invalidIdentifiers.length === depImport.identifiers.length) {
-        // If all identifiers are invalid, mark the entire import for removal
-        indexesToRemove.push({
-          startIndex: depImport.node.startIndex,
-          endIndex: depImport.node.endIndex,
-        });
-      } else {
-        // Remove only the invalid identifiers from the import statement
-        invalidIdentifiers.forEach((identifierNode) => {
-          indexesToRemove.push({
-            startIndex: identifierNode.startIndex,
-            endIndex: identifierNode.endIndex,
-          });
-        });
-      }
     });
 
     // Remove invalid imports and usages from the source code

packages/cli/src/languagesPlugins/python.test.ts

@@ -121,9 +121,7 @@ describe("Should get imports properly", () => {
 });
 
 describe("Should get exports properly", () => {
-  const entryPointPath = "api/index.py";
-
-  const plugin = new PythonPlugin(entryPointPath);
+  const plugin = new PythonPlugin("api/index.py");
 
   it.each([
     {
@@ -421,3 +419,113 @@ def bar(b):
     expect(exports[1].identifiers[0].identifierNode.text).toBe(`bar`);
   });
 });
+
+describe("Should cleanup unused imports", () => {
+  const entryPoint = "api/index.py";
+  const plugin = new PythonPlugin(entryPoint);
+
+  it.each([
+    {
+      sourceCode: `from module import foo`,
+      expectedSourceCode: ``,
+    },
+    {
+      sourceCode: `from module import foo as bar`,
+      expectedSourceCode: ``,
+    },
+    {
+      sourceCode: `from module import foo`,
+      expectedSourceCode: ``,
+    },
+    {
+      sourceCode: `from module import foo, bar as b`,
+      expectedSourceCode: ``,
+    },
+  ])(
+    "Should completely remove unused imports",
+    ({ sourceCode, expectedSourceCode }) => {
+      const updatedSourceCode = plugin.cleanupUnusedImports(
+        entryPoint,
+        sourceCode,
+      );
+
+      expect(updatedSourceCode).toBe(expectedSourceCode);
+    },
+  );
+
+  it("Should retain side effect import", () => {
+    const sourceCode = `import module`;
+
+    const updatedSourceCode = plugin.cleanupUnusedImports(
+      entryPoint,
+      sourceCode,
+    );
+
+    expect(updatedSourceCode).toBe(sourceCode);
+  });
+
+  it.each([
+    {
+      sourceCode: `
+from module import foo, bar
+bar()
+`,
+      expectedSourceCode: `
+from module import , bar
+bar()
+`,
+    },
+    {
+      sourceCode: `
+from module import foo, bar
+foo()
+    `,
+      expectedSourceCode: `
+from module import foo, 
+foo()
+    `,
+    },
+  ])(
+    "Should remove partially unused imports",
+    ({ sourceCode, expectedSourceCode }) => {
+      const updatedSourceCode = plugin.cleanupUnusedImports(
+        entryPoint,
+        sourceCode,
+      );
+
+      expect(updatedSourceCode).toBe(expectedSourceCode);
+    },
+  );
+  it.each([
+    {
+      sourceCode: `
+from module import foo as f, bar as b
+b()
+`,
+      expectedSourceCode: `
+from module import , bar as b
+b()
+`,
+    },
+    {
+      sourceCode: `
+from module import foo as f, bar as b
+f()
+`,
+      expectedSourceCode: `
+from module import foo as f, 
+f()
+`,
+    },
+  ])(
+    "Should remove partially unused imports with aliases",
+    ({ sourceCode, expectedSourceCode }) => {
+      const updatedSourceCode = plugin.cleanupUnusedImports(
+        entryPoint,
+        sourceCode,
+      );
+
+      expect(updatedSourceCode).toBe(expectedSourceCode);
+    },
+  );
+});