From 3dceb65e5731afe7631c2c15e687edf34b083277 Mon Sep 17 00:00:00 2001
From: "Daniel St. Jules" <danielst.jules@gmail.com>
Date: Tue, 15 Sep 2015 13:13:37 -0700
Subject: [PATCH] Add wrapStaticMethods and wrapInstanceMethods, as well as
 methodNames param

---
 README.md | 37 ++++++++++++++++++----------
 index.js  | 74 ++++++++++++++++++++++++++++++++++++++++++++++++-------
 test.js   |  4 +++
 3 files changed, 93 insertions(+), 22 deletions(-)

diff --git a/README.md b/README.md
index fc13d10..325548c 100644
--- a/README.md
+++ b/README.md
@@ -49,19 +49,7 @@ FakeDataStore.prototype.setAsync = Promise.coroutine(function*(key, value) {
 };
 ```
 
-That's where this library comes in. It offers a helpful little wrap function:
-
-``` javascript
-/**
- * Wraps all static and instance methods whose name ends with Async. Any
- * GeneratorFunction is wrapped with bluebird.coroutine(), and others with
- * bluebird.method().
- *
- * @param {function} klass The class to wrap
- */
-```
-
-Using it is simple:
+That's where this library comes in. Using it is simple:
 
 ``` javascript
 'use strict';
@@ -82,3 +70,26 @@ wrap(FakeDataStore);
 ```
 
 Clean ES6 classes and async methods!
+
+## API
+
+#### async-class.wrap(klass [, methodNames])
+
+Wraps all static and instance methods whose name ends with Async. Any
+GeneratorFunction is wrapped with bluebird.coroutine(), and others with
+bluebird.method(). Accepts an optional array of method names, wrapping
+only those found in the array, and disabling the Async suffix check.
+
+#### async-class.wrapStaticMethods(klass [, methodNames])
+
+Wraps all static methods whose name ends with Async. Any GeneratorFunction
+is wrapped with bluebird.coroutine(), and others with bluebird.method().
+Accepts an optional array of method names, wrapping only those found in the
+array, and disabling the Async suffix check.
+
+#### async-class.wrapInstanceMethods(klass [, methodNames])
+
+Wraps all instance methods whose name ends with Async. Any GeneratorFunction
+is wrapped with bluebird.coroutine(), and others with bluebird.method().
+Accepts an optional array of method names, wrapping only those found in the
+array, and disabling the Async suffix check.
diff --git a/index.js b/index.js
index a791b2f..416aae8 100644
--- a/index.js
+++ b/index.js
@@ -4,21 +4,71 @@ let Promise = require('bluebird');
 /**
  * Wraps all static and instance methods whose name ends with Async. Any
  * GeneratorFunction is wrapped with bluebird.coroutine(), and others with
- * bluebird.method().
+ * bluebird.method(). Accepts an optional array of method names, wrapping
+ * only those found in the array, and disabling the Async suffix check.
  *
- * @param {function} klass The class to wrap
+ * @param  {function} klass         The class to wrap
+ * @param  {string[]} [methodNames] Optional array of method names
+ * @throws {Error}    If methodNames is provided, but is not an array
  */
-function wrap(klass) {
-  [klass, klass.prototype].forEach(wrapFunctions);
+function wrap(klass, methodNames) {
+  validateMethodNames(methodNames);
+
+  wrapStaticMethods(klass, methodNames);
+  wrapInstanceMethods(klass, methodNames);
 }
 
-module.exports = {
-  wrap
-};
+/**
+ * Wraps all static methods whose name ends with Async. Any GeneratorFunction
+ * is wrapped with bluebird.coroutine(), and others with bluebird.method().
+ * Accepts an optional array of method names, wrapping only those found in the
+ * array, and disabling the Async suffix check.
+ *
+ * @param  {function} klass         The class to wrap
+ * @param  {string[]} [methodNames] Optional array of method names
+ * @throws {Error}    If methodNames is provided, but is not an array
+ */
+function wrapStaticMethods(klass, methodNames) {
+  validateMethodNames(methodNames);
+  wrapFunctions(klass, methodNames);
+}
+
+/**
+ * Wraps all instance methods whose name ends with Async. Any GeneratorFunction
+ * is wrapped with bluebird.coroutine(), and others with bluebird.method().
+ * Accepts an optional array of method names, wrapping only those found in the
+ * array, and disabling the Async suffix check.
+ *
+ * @param  {function} klass         The class to wrap
+ * @param  {string[]} [methodNames] Optional array of method names
+ * @throws {Error}    If methodNames is provided, but is not an array
+ */
+function wrapInstanceMethods(klass, methodNames) {
+  validateMethodNames(methodNames);
+  wrapFunctions(klass.prototype, methodNames);
+}
+
+/**
+ * Helper function that validates the methodNames parameter.
+ *
+ * @param {string[]} [methodNames] Optional array of method names
+ * @throws {Error}   If methodNames is provided, but is not an array
+ */
+function validateMethodNames(methodNames) {
+  if (methodNames && !(methodNames instanceof Array)) {
+    throw new Error('Optional methodNames should be an array if provided');
+  }
+}
 
-function wrapFunctions(target) {
+function wrapFunctions(target, methodNames) {
   Object.getOwnPropertyNames(target).forEach(function(key) {
-    if (!key.endsWith('Async') || typeof target[key] !== 'function') return;
+    if (methodNames) {
+      if (methodNames.indexOf(key) === -1) return;
+    } else if (!key.endsWith('Async')) {
+      return;
+    }
+
+    if (typeof target[key] !== 'function') return;
 
     if (target[key].constructor.name === 'GeneratorFunction') {
       target[key] = Promise.coroutine(target[key]);
@@ -27,3 +77,9 @@ function wrapFunctions(target) {
     }
   });
 }
+
+module.exports = {
+  wrap,
+  wrapStaticMethods,
+  wrapInstanceMethods
+};
diff --git a/test.js b/test.js
index e3a9704..18580a2 100644
--- a/test.js
+++ b/test.js
@@ -50,6 +50,10 @@ describe('async-class', function() {
   });
 
   describe('wrap', function() {
+    it('only modifies properties that are functions', function() {
+      expect(dataStore.store).to.be.instanceOf(Map)
+    });
+
     it('does not wrap instance methods that do not end with Async', function() {
       dataStore.store.set('foo', 'bar');
       let keys = dataStore.keys();