From 241cba6cc755bdcce0b972255368865e58c674a9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Eduardo=20Bou=C3=A7as?= <mail@eduardoboucas.com>
Date: Fri, 1 Mar 2024 09:59:14 +0000
Subject: [PATCH] feat: add `analysis.transformAST` hook

---
 readme.md                                    | 13 +++++++++++
 src/analyze.ts                               |  4 ++++
 src/node-file-trace.ts                       |  2 ++
 src/types.ts                                 |  2 ++
 test/unit.test.js                            | 23 ++++++++++++++++++--
 test/unit/imports-transform-ast/input.js     |  2 ++
 test/unit/imports-transform-ast/output.js    |  4 ++++
 test/unit/imports-transform-ast/package.json |  4 ++++
 8 files changed, 52 insertions(+), 2 deletions(-)
 create mode 100644 test/unit/imports-transform-ast/input.js
 create mode 100644 test/unit/imports-transform-ast/output.js
 create mode 100644 test/unit/imports-transform-ast/package.json

diff --git a/readme.md b/readme.md
index bb9218dc..414d26ee 100644
--- a/readme.md
+++ b/readme.md
@@ -172,6 +172,19 @@ const { fileList } = await nodeFileTrace(files, {
     computeFileReferences: true,
     // evaluate known bindings to assist with glob and file reference analysis
     evaluatePureExpressions: true,
+    // optional hook into the analysis step to inspect or modify the generated AST
+    transformAST: async (path, ast) => {
+      await walk(ast, {
+        async enter(node) {
+          if (
+            node.type === 'ImportDeclaration' &&
+            node.source.value === 'foo'
+          ) {
+            this.remove();
+          }
+        },
+      });
+    },
   },
 });
 ```
diff --git a/src/analyze.ts b/src/analyze.ts
index 39150d14..aa1252d9 100644
--- a/src/analyze.ts
+++ b/src/analyze.ts
@@ -338,6 +338,10 @@ export default async function analyze(
     }
   }
 
+  if (job.analysis.transformAST) {
+    await job.analysis.transformAST(id, ast);
+  }
+
   const importMetaUrl = pathToFileURL(id).href;
 
   const knownBindings: Record<
diff --git a/src/node-file-trace.ts b/src/node-file-trace.ts
index 80c01857..28c1ee6f 100644
--- a/src/node-file-trace.ts
+++ b/src/node-file-trace.ts
@@ -9,6 +9,7 @@ import analyze, { AnalyzeResult } from './analyze';
 import resolveDependency, { NotFoundError } from './resolve-dependency';
 import { isMatch } from 'micromatch';
 import { sharedLibEmit } from './utils/sharedlib-emit';
+import { Node } from './utils/types';
 import { join } from 'path';
 import { CachedFileSystem } from './fs';
 
@@ -61,6 +62,7 @@ export class Job {
     emitGlobs?: boolean;
     computeFileReferences?: boolean;
     evaluatePureExpressions?: boolean;
+    transformAST?: (path: string, node: Node) => Promise<void>;
   };
   private analysisCache: Map<string, AnalyzeResult>;
   public fileList: Set<string>;
diff --git a/src/types.ts b/src/types.ts
index 5557ff02..78a27b13 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -1,4 +1,5 @@
 import { Job } from './node-file-trace';
+import { Node } from './utils/types';
 
 export interface Stats {
   isFile(): boolean;
@@ -41,6 +42,7 @@ export interface NodeFileTraceOptions {
         emitGlobs?: boolean;
         computeFileReferences?: boolean;
         evaluatePureExpressions?: boolean;
+        transformAST?: (path: string, node: Node) => Promise<void>;
       };
   cache?: any;
   paths?: Record<string, string>;
diff --git a/test/unit.test.js b/test/unit.test.js
index c2f34426..39d6abf7 100644
--- a/test/unit.test.js
+++ b/test/unit.test.js
@@ -1,5 +1,6 @@
 const fs = require('fs');
 const { join, relative } = require('path');
+const { asyncWalk } = require('estree-walker');
 const { nodeFileTrace } = require('../out/node-file-trace');
 const gracefulFS = require('graceful-fs');
 const analyze = require('../out/analyze.js').default;
@@ -122,6 +123,25 @@ for (const { testName, isRoot } of unitTests) {
         inputFileNames.push('input-2.js', 'input-3.js', 'input-4.js');
       }
 
+      // disable analysis for basic-analysis unit tests
+      let analysis = !testName.startsWith('basic-analysis');
+
+      if (testName === 'imports-transform-ast') {
+        analysis = {
+          transformAST: async (path, ast) => {
+            expect(path).toEqual(join(unitPath, 'input.js'));
+
+            await asyncWalk(ast, {
+              async enter(node) {
+                if (node.type === 'ImportDeclaration') {
+                  this.remove();
+                }
+              },
+            });
+          },
+        };
+      }
+
       const { fileList, reasons } = await nodeFileTrace(
         inputFileNames.map((file) => join(unitPath, file)),
         {
@@ -135,8 +155,7 @@ for (const { testName, isRoot } of unitTests) {
           exportsOnly: testName.startsWith('exports-only'),
           ts: true,
           log: true,
-          // disable analysis for basic-analysis unit tests
-          analysis: !testName.startsWith('basic-analysis'),
+          analysis,
           mixedModules: true,
           // Ignore unit test output "actual.js", and ignore GitHub Actions preinstalled packages
           ignore: (str) =>
diff --git a/test/unit/imports-transform-ast/input.js b/test/unit/imports-transform-ast/input.js
new file mode 100644
index 00000000..ef2955e5
--- /dev/null
+++ b/test/unit/imports-transform-ast/input.js
@@ -0,0 +1,2 @@
+import { x } from '#x';
+console.log(x);
diff --git a/test/unit/imports-transform-ast/output.js b/test/unit/imports-transform-ast/output.js
new file mode 100644
index 00000000..abbbee67
--- /dev/null
+++ b/test/unit/imports-transform-ast/output.js
@@ -0,0 +1,4 @@
+[
+  "test/unit/imports-transform-ast/input.js",
+  "test/unit/imports-transform-ast/package.json"
+]
\ No newline at end of file
diff --git a/test/unit/imports-transform-ast/package.json b/test/unit/imports-transform-ast/package.json
new file mode 100644
index 00000000..cbf9e90c
--- /dev/null
+++ b/test/unit/imports-transform-ast/package.json
@@ -0,0 +1,4 @@
+{
+  "name": "x",
+  "type": "module"
+}