diff options
Diffstat (limited to 'lib')
-rw-r--r-- | lib/__with__.js | 15 |
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(); + } } }; } |