From 12a2d07c215e427698179eb6ef29cb5a7f31b5ce Mon Sep 17 00:00:00 2001
From: Daniel Tschinder <daniel@tschinder.de>
Date: Thu, 8 Feb 2018 19:06:29 +0100
Subject: [PATCH] feat: run plugin also on Program exit to handle dynamically
 added imports from other transforms (#269)

---
 src/index.js               | 10 ++++++++++
 src/transformers/call.js   |  5 +++++
 src/transformers/import.js |  5 +++++
 test/dynamicImport.test.js | 22 ++++++++++++++++++++++
 4 files changed, 42 insertions(+)

diff --git a/src/index.js b/src/index.js
index e52af73..1f8913c 100644
--- a/src/index.js
+++ b/src/index.js
@@ -18,6 +18,9 @@ const visitor = {
     enter(programPath, state) {
       programPath.traverse(importVisitors, state);
     },
+    exit(programPath, state) {
+      programPath.traverse(importVisitors, state);
+    },
   },
 };
 
@@ -29,7 +32,14 @@ export default ({ types }) => ({
 
     const currentFile = file.opts.filename;
     this.normalizedOpts = normalizeOptions(currentFile, this.opts);
+    // We need to keep track of all handled nodes so we do not try to transform them twice,
+    // because we run before (enter) and after (exit) all nodes are handled
+    this.moduleResolverVisited = new Set();
   },
 
   visitor,
+
+  post() {
+    this.moduleResolverVisited.clear();
+  },
 });
diff --git a/src/transformers/call.js b/src/transformers/call.js
index 8709192..f5d463e 100644
--- a/src/transformers/call.js
+++ b/src/transformers/call.js
@@ -6,12 +6,17 @@ import {
 
 
 export default function transformCall(nodePath, state) {
+  if (state.moduleResolverVisited.has(nodePath)) {
+    return;
+  }
+
   const calleePath = nodePath.get('callee');
   const isNormalCall = state.normalizedOpts.transformFunctions.some(
     pattern => matchesPattern(state.types, calleePath, pattern),
   );
 
   if (isNormalCall || isImportCall(state.types, nodePath)) {
+    state.moduleResolverVisited.add(nodePath);
     mapPathString(nodePath.get('arguments.0'), state);
   }
 }
diff --git a/src/transformers/import.js b/src/transformers/import.js
index 9dc78ae..c2970b6 100644
--- a/src/transformers/import.js
+++ b/src/transformers/import.js
@@ -2,5 +2,10 @@ import { mapPathString } from '../utils';
 
 
 export default function transformImport(nodePath, state) {
+  if (state.moduleResolverVisited.has(nodePath)) {
+    return;
+  }
+  state.moduleResolverVisited.add(nodePath);
+
   mapPathString(nodePath.get('source'), state);
 }
diff --git a/test/dynamicImport.test.js b/test/dynamicImport.test.js
index 7b7b6c4..cb60bcd 100644
--- a/test/dynamicImport.test.js
+++ b/test/dynamicImport.test.js
@@ -56,4 +56,26 @@ describe('import()', () => {
 
     expect(result.code).toBe('import("").then(() => {}).catch(() => {});');
   });
+
+  it('should handle imports added by other transforms', () => {
+    const options = {
+      ...transformerOpts,
+      plugins: [
+        function fakePlugin({ types }) {
+          return {
+            visitor: {
+              Identifier(path) {
+                path.replaceWith(types.Import());
+              },
+            },
+          };
+        },
+        ...transformerOpts.plugins,
+      ],
+    };
+    const code = 'boo("components/Header/SubHeader");';
+    const result = transform(code, options);
+
+    expect(result.code).toBe('import("./test/testproject/src/components/Header/SubHeader");');
+  });
 });