1 files changed, 14 insertions, 3 deletions
diff --git a/lib/retryable.js b/lib/retryable.js
index 9422eed..ad99801 100644
--- a/lib/retryable.js
+++ b/lib/retryable.js
@@ -1,6 +1,7 @@
 import retry from './retry';
 import initialParams from './internal/initialParams';
-import wrapAsync from './internal/wrapAsync';
+import {default as wrapAsync, isAsync} from './internal/wrapAsync';
+import { promiseCallback, PROMISE_SYMBOL } from './internal/promiseCallback';
 
 /**
  * A close relative of [`retry`]{@link module:ControlFlow.retry}.  This method
@@ -14,7 +15,8 @@ import wrapAsync from './internal/wrapAsync';
  * @see [async.retry]{@link module:ControlFlow.retry}
  * @category Control Flow
  * @param {Object|number} [opts = {times: 5, interval: 0}| 5] - optional
- * options, exactly the same as from `retry`
+ * options, exactly the same as from `retry`, except for a `opts.arity` that
+ * is the arity of the `task` function, defaulting to `task.length`
  * @param {AsyncFunction} task - the asynchronous function to wrap.
  * This function will be passed any arguments passed to the returned wrapper.
  * Invoked with (...args, callback).
@@ -30,13 +32,21 @@ import wrapAsync from './internal/wrapAsync';
  *     })]
  * }, callback);
  */
-export default function (opts, task) {
+export default function retryable (opts, task) {
     if (!task) {
         task = opts;
         opts = null;
     }
+    let arity = (opts && opts.arity) || task.length
+    if (isAsync(task)) {
+        arity += 1
+    }
     var _task = wrapAsync(task);
     return initialParams((args, callback) => {
+        if (args.length < arity - 1 || callback == null) {
+            args.push(callback)
+            callback = promiseCallback()
+        }
         function taskFn(cb) {
             _task(...args, cb);
         }
@@ -44,5 +54,6 @@ export default function (opts, task) {
         if (opts) retry(opts, taskFn, callback);
         else retry(taskFn, callback);
 
+        return callback[PROMISE_SYMBOL]
     });
 }