Add promise feature to __with__

1 files changed, 12 insertions, 3 deletions

diff --git a/lib/__with__.js b/lib/__with__.js
index ef29c8c..387714d 100644
--- a/lib/__with__.js
+++ b/lib/__with__.js
@@ -15,7 +15,9 @@ function __with__() {
     var args = arguments;
 
     return function (callback) {
-        var undo;
+        var undo,
+            returned,
+            isPromise;
 
         if (typeof callback !== "function") {
             throw new TypeError("__with__ expects a callback function");
@@ -24,9 +26,16 @@ function __with__() {
         undo = module.exports.__set__.apply(null, args);
 
         try {
-            callback();
+            returned = callback();
+            isPromise = returned && typeof returned.then === "function";
+            if (isPromise) {
+                returned.then(undo, undo);
+                return returned;
+            }
         } finally {
-            undo();
+            if (!isPromise) {
+                undo();
+            }
         }
     };
 }