1 files changed, 29 insertions, 2 deletions

diff --git a/lib/moduleEnv.js b/lib/moduleEnv.js
index 8ad7891..789ea2f 100644
--- a/lib/moduleEnv.js
+++ b/lib/moduleEnv.js
@@ -2,12 +2,15 @@
 
 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;
 
@@ -26,7 +29,6 @@ function load(targetModule) {
 function reset() {
     Module.wrapper[0] = moduleWrapper0;
     Module.wrapper[1] = moduleWrapper1;
-    restoreExtensions();
 }
 
 function inject(prelude, appendix) {
@@ -47,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