Refactor code to match module style

1 files changed, 24 insertions, 21 deletions

diff --git a/lib/__set__.js b/lib/__set__.js
index ad0e5b0..feea24a 100644
--- a/lib/__set__.js
+++ b/lib/__set__.js
@@ -5,18 +5,17 @@
  * All variables within this function are namespaced in the arguments array because every
  * var declaration could possibly clash with a variable in the module scope.
  *
- * @param {!String|!Object} varName name of the variable to set
+ * @param {String|Object} varName name of the variable to set
  * @param {String} varValue new value
  * @throws {TypeError}
  * @throws {ReferenceError} When the variable is unknown
  * @return {*}
  */
-
 function __set__() {
     arguments.varName = arguments[0];
     arguments.varValue = arguments[1];
     arguments.src = "";
-    var snapshot = {};
+    arguments.snapshot = {};
 
     if (typeof arguments[0] === "object" && arguments.length === 1) {
         arguments.env = arguments.varName;
@@ -27,7 +26,7 @@ function __set__() {
             if (arguments.env.hasOwnProperty(arguments.varName)) {
                 arguments.varValue = arguments.env[arguments.varName];
                 arguments.src += arguments.varName + " = arguments.env." + arguments.varName + "; ";
-                snapshot[arguments.varName] = eval(arguments.varName);
+                arguments.snapshot[arguments.varName] = eval(arguments.varName);
             }
         }
     } else if (typeof arguments.varName === "string" && arguments.length === 2) {
@@ -35,32 +34,36 @@ function __set__() {
             throw new TypeError("__set__ expects a non-empty string as a variable name");
         }
         arguments.src = arguments.varName + " = arguments.varValue;";
-        snapshot[arguments.varName] = eval(arguments.varName);
+        arguments.snapshot[arguments.varName] = eval(arguments.varName);
     } else {
         throw new TypeError("__set__ expects an environment object or a non-empty string as a variable name");
     }
 
     eval(arguments.src);
-    return function() {
-      __set__(snapshot);
-    };
+
+    return function (snapshot) {
+        __set__(snapshot);
+    }.bind(null, arguments.snapshot);
 }
 
 function __with__() {
-  var args = arguments;
-  return function(callback) {
-    if (typeof callback !== "function") {
-      throw new TypeError("__with__ expects a callback function")
-    }
+    var args = arguments;
 
-    var undo = __set__.apply(null, args)
-    try {
-      callback();
-    }
-    finally {
-      undo();
+    return function (callback) {
+        var undo;
+
+        if (typeof callback !== "function") {
+            throw new TypeError("__with__ expects a callback function")
+        }
+
+        undo = __set__.apply(null, args);
+
+        try {
+            callback();
+        } finally {
+            undo();
+        }
     }
-  }
 }
 
-module.exports = {"__set__": __set__, "__with__": __with__}
+module.exports = {"__set__": __set__, "__with__": __with__};