Refactor code

2 files changed, 40 insertions, 22 deletions

diff --git a/lib/moduleEnv.js b/lib/moduleEnv.js
index dd003d7..09bfd39 100644
--- a/lib/moduleEnv.js
+++ b/lib/moduleEnv.js
@@ -3,12 +3,14 @@
 var Module = require("module"),
     fs = require("fs"),
     babelCore = require("babel-core"),
+    // Requiring the babel plugin here because otherwise it will be lazy-loaded by Babel during rewire()
+    transformBlockScoping = require("babel-plugin-transform-es2015-block-scoping"),
     coffee;
 
-    // caching original wrapper
 var moduleWrapper0 = Module.wrapper[0],
     moduleWrapper1 = Module.wrapper[1],
     originalExtensions = {},
+    matchCoffeeExt = /\.coffee$/,
     nodeRequire,
     currentModule;
 
@@ -18,32 +20,30 @@ function load(targetModule) {
     currentModule = targetModule;
 
     registerExtensions();
-
     targetModule.load(targetModule.id);
 
     // This is only necessary if nothing has been required within the module
     reset();
 }
 
-function compile(targetModule, src, filename) {
-    nodeRequire = targetModule.require;
-    targetModule.require = requireProxy;
-    targetModule.filename = filename;
-    currentModule = targetModule;
+function compileJavaScript(src, targetPath) {
+    return babelCore.transform(stripBOM(src), {
+        plugins: [transformBlockScoping],
+        retainLines: true,
+        filename: targetPath
+    }).code;
+}
 
-    var convertedSrc = babelCore.transform(stripBOM(src), {
-        plugins: ["transform-es2015-constants"]
+function compileCoffeeScript(src, targetPath) {
+    return coffee.compile(stripBOM(src), {
+        filename: targetPath,
+        bare: true
     });
-
-    targetModule._compile(convertedSrc.code, filename);
-
-    reset();
 }
 
 function reset() {
     Module.wrapper[0] = moduleWrapper0;
     Module.wrapper[1] = moduleWrapper1;
-    restoreExtensions();
 }
 
 function inject(prelude, appendix) {
@@ -64,23 +64,48 @@ function requireProxy(path) {
 }
 
 function registerExtensions() {
+    var originalJsExtension = require.extensions[".js"];
     var originalCoffeeExtension = require.extensions[".coffee"];
 
+    if (originalJsExtension) {
+        originalExtensions.js = originalJsExtension;
+    }
     if (originalCoffeeExtension) {
         originalExtensions.coffee = originalCoffeeExtension;
     }
+    require.extensions[".js"] = jsExtension;
     require.extensions[".coffee"] = coffeeExtension;
 }
 
 function restoreExtensions() {
+    if ("js" in originalExtensions) {
+        require.extensions[".js"] = originalExtensions.js;
+    }
     if ("coffee" in originalExtensions) {
         require.extensions[".coffee"] = originalExtensions.coffee;
     }
 }
 
+function jsExtension(module, filename) {
+    var _compile = module._compile;
+
+    module._compile = function (content, filename) {
+        content = babelCore.transform(content, {
+            plugins: ["transform-es2015-block-scoping"],
+            retainLines: true,
+            filename: filename
+        }).code;
+        _compile.call(module, content, filename);
+    };
+
+    restoreExtensions();
+    originalExtensions.js(module, filename);
+}
+
 function coffeeExtension(module, filename) {
     var content = stripBOM(fs.readFileSync(filename, "utf8"));
 
+    restoreExtensions();
     content = coffee.compile(content, {
         filename: filename,
         bare: true
@@ -108,5 +133,4 @@ try {
 }
 
 exports.load = load;
-exports.compile = compile;
 exports.inject = inject;
diff --git a/lib/rewire.js b/lib/rewire.js
index 536929d..cbd24c2 100644
--- a/lib/rewire.js
+++ b/lib/rewire.js
@@ -55,13 +55,7 @@ function internalRewire(parentModulePath, targetPath) {
     }
 
     moduleEnv.inject(prelude, appendix);
-
-    if (targetPath.match(/\.coffee$/)) {
-        moduleEnv.load(targetModule);
-    } else {
-        moduleEnv.compile(targetModule, src, targetPath);
-    }
-
+    moduleEnv.load(targetModule);
 
     return targetModule.exports;
 }