Allow for custom node type keys (#139)

Author: nzakas

esquery.js

@@ -117,7 +117,10 @@ function generateMatcher(selector) {
 
         case 'identifier': {
             const value = selector.value.toLowerCase();
-            return (node) => value === node.type.toLowerCase();
+            return (node, ancestry, options) => {
+                const nodeTypeKey = (options && options.nodeTypeKey) || 'type';
+                return value === node[nodeTypeKey].toLowerCase();
+            };
         }
 
         case 'field': {
@@ -294,8 +297,12 @@ function generateMatcher(selector) {
         }
 
         case 'class': {
+
             const name = selector.name.toLowerCase();
-            return (node, ancestry) => {
+            return (node, ancestry, options) => {
+
+                if (options && options.nodeTypeKey) return false;
+
                 switch(name){
                     case 'statement':
                         if(node.type.slice(-9) === 'Statement') return true;
@@ -333,6 +340,7 @@ function generateMatcher(selector) {
  */
 /**
  * @typedef {object} ESQueryOptions
+ * @property {string} [nodeTypeKey="type"] By passing `nodeTypeKey`, we can allow other ASTs to use ESQuery.
  * @property { { [nodeType: string]: string[] } } [visitorKeys] By passing `visitorKeys` mapping, we can extend the properties of the nodes that traverse the node.
  * @property {TraverseOptionFallback} [fallback] By passing `fallback` option, we can control the properties of traversing nodes when encountering unknown nodes.
  */
@@ -363,7 +371,9 @@ function matches(node, selector, ancestry, options) {
  * @returns {string[]} Visitor keys of the node.
  */
 function getVisitorKeys(node, options) {
-    const nodeType = node.type;
+    const nodeTypeKey = (options && options.nodeTypeKey) || 'type';
+
+    const nodeType = node[nodeTypeKey];
     if (options && options.visitorKeys && options.visitorKeys[nodeType]) {
         return options.visitorKeys[nodeType];
     }
@@ -375,18 +385,20 @@ function getVisitorKeys(node, options) {
     }
     // 'iteration' fallback
     return Object.keys(node).filter(function (key) {
-        return key !== 'type';
+        return key !== nodeTypeKey;
     });
 }
 
 
 /**
  * Check whether the given value is an ASTNode or not.
  * @param {any} node The value to check.
+ * @param {ESQueryOptions|undefined} options The options to use.
  * @returns {boolean} `true` if the value is an ASTNode.
  */
-function isNode(node) {
-    return node !== null && typeof node === 'object' && typeof node.type === 'string';
+function isNode(node, options) {
+    const nodeTypeKey = (options && options.nodeTypeKey) || 'type';
+    return node !== null && typeof node === 'object' && typeof node[nodeTypeKey] === 'string';
 }
 
 /**
@@ -417,7 +429,7 @@ function sibling(node, matcher, ancestry, side, options) {
                 upperBound = listProp.length;
             }
             for (let k = lowerBound; k < upperBound; ++k) {
-                if (isNode(listProp[k]) && matcher(listProp[k], ancestry, options)) {
+                if (isNode(listProp[k], options) && matcher(listProp[k], ancestry, options)) {
                     return true;
                 }
             }
@@ -445,10 +457,10 @@ function adjacent(node, matcher, ancestry, side, options) {
         if (Array.isArray(listProp)) {
             const idx = listProp.indexOf(node);
             if (idx < 0) { continue; }
-            if (side === LEFT_SIDE && idx > 0 && isNode(listProp[idx - 1]) && matcher(listProp[idx - 1], ancestry, options)) {
+            if (side === LEFT_SIDE && idx > 0 && isNode(listProp[idx - 1], options) && matcher(listProp[idx - 1], ancestry, options)) {
                 return true;
             }
-            if (side === RIGHT_SIDE && idx < listProp.length - 1 && isNode(listProp[idx + 1]) &&  matcher(listProp[idx + 1], ancestry, options)) {
+            if (side === RIGHT_SIDE && idx < listProp.length - 1 && isNode(listProp[idx + 1], options) &&  matcher(listProp[idx + 1], ancestry, options)) {
                 return true;
             }
         }

tests/fixtures/customNodesWithKind.js

@@ -0,0 +1,38 @@
+export default {
+    kind: 'CustomRoot',
+    list: [
+        {
+            kind: 'CustomChild',
+            name: 'one',
+            sublist: [{ kind: 'CustomGrandChild' }],
+        },
+        {
+            kind: 'CustomChild',
+            name: 'two',
+            sublist: [],
+        },
+        {
+            kind: 'CustomChild',
+            name: 'three',
+            sublist: [
+                { kind: 'CustomGrandChild' },
+                { kind: 'CustomGrandChild' },
+            ],
+        },
+        {
+            kind: 'CustomChild',
+            name: 'four',
+            sublist: [
+                { kind: 'CustomGrandChild' },
+                { kind: 'CustomGrandChild' },
+                { kind: 'CustomGrandChild' },
+            ],
+        },
+        {
+            kind: 'CustomExpression'
+        },
+        {
+            kind: 'CustomStatement'
+        }
+    ],
+};

tests/matches.js

@@ -3,6 +3,7 @@ import forLoop from './fixtures/forLoop.js';
 import simpleProgram from './fixtures/simpleProgram.js';
 import conditional from './fixtures/conditional.js';
 import customNodes from './fixtures/customNodes.js';
+import customNodesWithKind from './fixtures/customNodesWithKind.js';
 
 describe('matches', function () {
     it('falsey node', function () {
@@ -167,6 +168,41 @@ describe('matches with custom AST and custom visitor keys', function () {
     });
 });
 
+describe('matches with custom AST and nodeTypeKey and custom visitor keys', function () {
+    it('adjacent/sibling', function () {
+        const options = {
+            visitorKeys: {
+                CustomRoot: ['list'],
+                CustomChild: ['sublist'],
+                CustomGrandChild: [],
+                CustomStatement: [],
+                CustomExpression: []
+            },
+            nodeTypeKey: 'kind'
+        };
+
+        let selector = esquery.parse('CustomChild + CustomChild');
+        assert.doesNotThrow(() => {
+            esquery.matches(
+                customNodesWithKind.list[1],
+                selector,
+                [customNodesWithKind],
+                options
+            );
+        });
+
+        selector = esquery.parse('CustomChild ~ CustomChild');
+        assert.doesNotThrow(() => {
+            esquery.matches(
+                customNodesWithKind.list[1],
+                selector,
+                [customNodesWithKind],
+                options
+            );
+        });
+    });
+});
+
 describe('matches with custom AST and fallback option', function () {
     it('adjacent/sibling', function () {
         const options = {

tests/queryClass.js

@@ -1,5 +1,6 @@
 import esquery from '../esquery.js';
 import ast from './fixtures/allClasses.js';
+import customNodesWithKind from './fixtures/customNodesWithKind.js';
 
 describe('Class query', function () {
 
@@ -66,4 +67,24 @@ describe('Class query', function () {
         assert.equal(10, matches.length);
     });
 
+    it('custom nodes with :expression, :statement', function () {
+        const options = {
+            visitorKeys: {
+                CustomRoot: ['list'],
+                CustomChild: ['sublist'],
+                CustomGrandChild: [],
+                CustomStatement: [],
+                CustomExpression: []
+            },
+            nodeTypeKey: 'kind'
+        };
+
+        const matches1 = esquery(customNodesWithKind, ':expression', options);
+        assert.equal(0, matches1.length);
+
+        const matches2 = esquery(customNodesWithKind, ':statement', options);
+        assert.equal(0, matches2.length);
+    });
+
+
 });