Add support for visiting nodes of multiple types

Closes GH-10.

Author: gibson042

index.js

@@ -3,12 +3,14 @@
 /* Expose. */
 module.exports = visit;
 
+var is = require('unist-util-is');
+
 /* Visit. */
-function visit(tree, type, visitor, reverse) {
-  if (typeof type === 'function') {
+function visit(tree, test, visitor, reverse) {
+  if (typeof test === 'function' && typeof visitor !== 'function') {
     reverse = visitor;
-    visitor = type;
-    type = null;
+    visitor = test;
+    test = null;
   }
 
   one(tree);
@@ -19,7 +21,7 @@ function visit(tree, type, visitor, reverse) {
 
     index = index || (parent ? 0 : null);
 
-    if (!type || node.type === type) {
+    if (!test || node.type === test || is(test, node, index, parent || null)) {
       result = visitor(node, index, parent || null);
     }
 

package.json

@@ -25,6 +25,9 @@
   "files": [
     "index.js"
   ],
+  "dependencies": {
+    "unist-util-is": "^2.1.1"
+  },
   "devDependencies": {
     "browserify": "^14.0.0",
     "esmangle": "^1.0.0",

readme.md

@@ -39,16 +39,16 @@ Yields:
 
 ## API
 
-### `visit(node[, type], visitor[, reverse])`
+### `visit(node[, test], visitor[, reverse])`
 
-Visit nodes.  Optionally by node type.  Optionally in reverse.
+Visit nodes.  Optionally filtering nodes.  Optionally in reverse.
 
 ###### Parameters
 
 *   `node` ([`Node`][node])
     — Node to search
-*   `type` (`string`, optional)
-    — Node type
+*   `test` ([`Test`][is], optional)
+    — Node type or other [`is`][is]-compatible test
 *   `visitor` ([Function][visitor])
     — Visitor invoked when a node is found
 *   `reverse` (`boolean`, default: `false`)
@@ -59,7 +59,7 @@ Visit nodes.  Optionally by node type.  Optionally in reverse.
 
 #### `stop? = visitor(node, index, parent)`
 
-Invoked when a node (when `type` is given, matching `type`) is found.
+Invoked when a node (matching `test`, if given) is found.
 
 ###### Parameters
 
@@ -114,4 +114,6 @@ Invoked when a node (when `type` is given, matching `type`) is found.
 
 [node]: https://github.com/syntax-tree/unist#node
 
+[is]: https://github.com/syntax-tree/unist-util-is#istest-node-index-parent-context
+
 [visitor]: #stop--visitornode-index-parent

test.js

@@ -9,6 +9,7 @@ var tree = remark().parse('Some _emphasis_, **importance**, and `code`.');
 var STOP = 5;
 
 var textNodes = 6;
+var codeNodes = 1;
 
 var types = [
   'root',
@@ -78,15 +79,69 @@ test('unist-util-visit', function (t) {
     st.end();
   });
 
-  t.test('should only visit given `types`', function (st) {
+  t.test('should only visit a given `type`', function (st) {
     var n = 0;
 
     visit(tree, 'text', function (node) {
       n++;
       st.equal(node.type, 'text');
     });
 
-    st.equal(n, textNodes, 'should visit all nodes');
+    st.equal(n, textNodes, 'should visit all matching nodes');
+
+    st.end();
+  });
+
+  t.test('should only visit given `type`s', function (st) {
+    var n = 0;
+    var types = ['text', 'inlineCode'];
+
+    visit(tree, types, function (node) {
+      n++;
+      st.ok(types.indexOf(node.type) !== -1, 'should be a requested type: ' + node.type);
+    });
+
+    st.equal(n, textNodes + codeNodes, 'should visit all matching nodes');
+
+    st.end();
+  });
+
+  t.test('should accept any `is`-compatible test', function (st) {
+    var n = 0;
+    var test = function (node, index) {
+      return index > 3;
+    };
+
+    visit(tree, test, function (node, index, parent) {
+      n++;
+      var parentType = parent && parent.type;
+      st.ok(index > 3, 'should be a requested node: ' + parentType + '/[' + index + ']');
+    });
+
+    st.equal(n, 3, 'should visit all matching nodes');
+
+    st.end();
+  });
+
+  t.test('should accept an array of `is`-compatible tests', function (st) {
+    var n = 0;
+    var tests = [
+      function (node) {
+        return node.type === 'root';
+      },
+      'paragraph',
+      {value: '.'},
+      ['emphasis', 'strong']
+    ];
+    var expectedTypes = ['root', 'paragraph', 'emphasis', 'strong'];
+
+    visit(tree, tests, function (node) {
+      n++;
+      st.ok(expectedTypes.indexOf(node.type) !== -1 || node.value === '.',
+          'should be a requested type: ' + node.type);
+    });
+
+    st.equal(n, 5, 'should visit all matching nodes');
 
     st.end();
   });