1 files changed, 29 insertions, 20 deletions

diff --git a/lib/moduleEnv.js b/lib/moduleEnv.js
index 29c407c..789ea2f 100644
--- a/lib/moduleEnv.js
+++ b/lib/moduleEnv.js
@@ -3,47 +3,32 @@
 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;
 
-function load(targetModule, isTransform) {
+function load(targetModule) {
     nodeRequire = targetModule.require;
     targetModule.require = requireProxy;
     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;
-
-    var convertedSrc = babelCore.transform(stripBOM(src), {
-        plugins: ["transform-es2015-constants"]
-    });
-
-    targetModule._compile(convertedSrc.code, filename);
-
-    reset();
-}
-
 function reset() {
     Module.wrapper[0] = moduleWrapper0;
     Module.wrapper[1] = moduleWrapper1;
-    restoreExtensions();
 }
 
 function inject(prelude, appendix) {
@@ -64,23 +49,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 +118,4 @@ try {
 }
 
 exports.load = load;
-exports.compile = compile;
 exports.inject = inject;