jscodeshift 那点秘密

在重构利器 jscodeshift 的末尾，留下了几个问题：

怎么做到同时满足 JavaScript 和 TypeScript 的解析？
jscodeshift 是如何实现链式调用的？
registerMethods 怎么实现，怎么使用？
测试工具 testUtils 做了哪些封装？

上一篇文章中我们举了一个 🌰：

export const sum = (a, b) => {
  console.log('计算下面两个数的和：', a, b);
  return a + b;
};

export const minus = (a, b) => {
  console.log('计算下面两个数的差：' + a + ',' + b);
  return a - b;
};

export const multiply = (a, b) => {
  console.warn('计算下面两个数的积：', a, b);
  return a * b;
};

export const divide = (a, b) => {
  console.error(`计算下面两个数相除 ${a}, ${b}`);
  return a / b;
};

转换器 transform 如下：

module.exports = (fileInfo, api, options) => {
  const j = api.jscodeshift;
  
  return j(fileInfo.source)
    .find(j.ExpressionStatement, {
      expression: {
        callee: {
          type: 'MemberExpression',
          object: { type: 'Identifier', name: 'console' },
        }
      },
    })
    .remove()
    .toSource();
};

咱们从入口开始分析，先看看这个 transform 文件的几个参数是什么？这部分代码位于 src/Worker.js：

function run(data) {
  const files = data.files;
  const options = data.options || {};
  if (!files.length) {
    finish();
    return;
  }
  async.each(
    files,
    function(file, callback) {
      fs.readFile(file, async function(err, source) {
        if (err) {
          updateStatus('error', file, 'File error: ' + err);
          callback();
          return;
        }
        source = source.toString();
        try {
          // 获取解析器 🛠️
          const jscodeshift = prepareJscodeshift(options);
          
          // 这里的transform就是我们写在transform.js中的函数，对应的参数就是从这里来的
          const out = await transform(
            // 上述 🌰 中的 fileInfo 参数
            {
              path: file,
              source: source,
            },
            // 上述 🌰 中的 api 参数
            {
              j: jscodeshift,
              jscodeshift: jscodeshift,
              stats: options.dry ? stats : empty,
              report: msg => report(file, msg),
            },
            // 上述 🌰 中的 options 参数
            options
          );
          // 无结果或者跟源码一样时
          if (!out || out === source) {
            updateStatus(out ? 'nochange' : 'skip', file);
            callback();
            return;
          }
          // --print 控制台打印结果
          if (options.print) {
            console.log(out); // eslint-disable-line no-console
          }
          
          // --dry 🤪🤪 如果不是陪跑，会将结果写入文件
          if (!options.dry) {
            writeFileAtomic(file, out, function(err) {
              if (err) {
                updateStatus('error', file, 'File writer error: ' + err);
              } else {
                updateStatus('ok', file);
              }
              callback();
            });
          } else {
            updateStatus('ok', file);
            callback();
          }
        } catch(err) {
          updateStatus(
            'error',
            file,
            'Transformation error ('+ err.message.replace(/\n/g, ' ') + ')\n' + trimStackTrace(err.stack)
          );
          callback();
        }
      });
    },
    function(err) {
      if (err) {
        updateStatus('error', '', 'This should never be shown!');
      }
      free();
    }
  );
}

run 的逻辑很简单，遍历文件调用 transform 函数去转换，同时将处理状态通过调用 notify 同步到父进程。先聚焦在获取解析器上：

/**
 * 获取 jscodeshift 解析器
 */
function prepareJscodeshift(options) {
  const parser = parserFromTransform ||
    getParser(options.parser, options.parserConfig);
  return jscodeshift.withParser(parser);
}

上面代码中，parserFromTransform 的定义在 setup ：

function setup(tr, babel) {
  // 默认情况下，也会默认设置一些预设和插件
  if (babel === 'babel') {
    const presets = [];
    if (presetEnv) {
      presets.push([
        presetEnv.default,
        {targets: {node: true}},
      ]);
    }
    presets.push(
      /\.tsx?$/.test(tr) ?
        require('@babel/preset-typescript').default :
        require('@babel/preset-flow').default
    );

    require('@babel/register')({
      babelrc: false,
      presets,
      plugins: [
        require('@babel/plugin-proposal-class-properties').default,
        require('@babel/plugin-proposal-nullish-coalescing-operator').default,
        require('@babel/plugin-proposal-optional-chaining').default,
        require('@babel/plugin-transform-modules-commonjs').default,
      ],
      extensions: [...DEFAULT_EXTENSIONS, '.ts', '.tsx'],
      // By default, babel register only compiles things inside the current working directory.
      // https://github.com/babel/babel/blob/2a4f16236656178e84b05b8915aab9261c55782c/packages/babel-register/src/node.js#L140-L157
      ignore: [
        // Ignore parser related files
        /@babel\/parser/,
        /\/flow-parser\//,
        /\/recast\//,
        /\/ast-types\//,
      ],
    });
  }

  // 获取transform函数
  const module = require(tr);
  // 兼容 ES module default 写法
  transform = typeof module.default === 'function' ?
    module.default :
    module;
  
  // 在 transform 有导出 parser 这个变量，这里就会去获取对应的解析器，可传的值有 babylon、flow、ts、tsx、babel
  if (module.parser) {
    parserFromTransform = typeof module.parser === 'string' ?
      getParser(module.parser) :
      module.parser;
  }
}

// src/getParser.js
'use strict';

module.exports = function getParser(parserName, options) {
  switch (parserName) {
    case 'babylon':
      return require('../parser/babylon')(options);
    case 'flow':
      return require('../parser/flow')(options);
    case 'ts':
      return require('../parser/ts')(options);
    case 'tsx':
      return require('../parser/tsx')(options);
    case 'babel':
    default:
      return require('../parser/babel5Compat')(options);
  }
};

setup 在 src/Runner.js 中 require('Worker.js') 时被调用，从命令行参数中可以看到默认是使用 babel 作为解析器。然后注册默认的 presets 和 plugins。最后判断用户是否自定义了 parser，是的话通过 getParser 去获取解析配置。
咱们来看下 parserName = ‘ts’ 的配置：

'use strict';

const babylon = require('@babel/parser');
const options = require('./tsOptions');

/**
 * Doesn't accept custom options because babylon should be used directly in
 * that case.
 */
module.exports = function() {
  return {
    parse(code) {
      return babylon.parse(code, options);
    },
  };
};

// ./tsOptions.js
'use strict';

/**
 * Options shared by the TypeScript and TSX parsers.
 */
module.exports = {
  sourceType: 'module',
  // 允许代码任何地方进行 import、export
  allowImportExportEverywhere: true,
  // 允许在函数外面进行 return
  allowReturnOutsideFunction: true,
  // 从第一行开始解析代码
  startLine: 1,
  // 将所有解析的标记添加到文件节点上的标记属性
  tokens: true,
  // https://babeljs.io/docs/en/babel-parser#plugins
  plugins: [
    'asyncGenerators',
    'bigInt', // 
    'classPrivateMethods',
    'classPrivateProperties',
    'classProperties',
    'decorators-legacy',
    'doExpressions',
    'dynamicImport',
    'exportDefaultFrom',
    'exportExtensions',
    'exportNamespaceFrom',
    'functionBind',
    'functionSent',
    'importMeta',
    'nullishCoalescingOperator',
    'numericSeparator',
    'objectRestSpread',
    'optionalCatchBinding',
    'optionalChaining',
    ['pipelineOperator', { proposal: 'minimal' }],
    'throwExpressions',
    'typescript'
  ],
};

可以看到解析 ts 是使用了 @babel/plugin-transform-typescript 插件。了解完通过 getParser 获取解析对象的配置后，接着看 prepareJscodeshift 中 withParser 方法是什么时候挂到 jscodeshift 上？还有哪些 API 供我们使用？注册逻辑定义在 src\core.js ：

/**
 * 编译主入口，支持各种不同的参数类型
 *
 * - 源码字符串（用recast去编译）
 * - 一个ast节点
 * - 一个node路径
 * - 一系列ast节点集合
 * - 一系列node路径集合
 *
 * @exports jscodeshift
 * @param {Node|NodePath|Array|string} source
 * @param {Object} options 传给recast做编译的参数
 * @return {Collection}
 */
function core(source, options) {
  return typeof source === 'string' ?
    fromSource(source, options) :
    fromAST(source);
}

/**
 * 返回指定了编译器的编译函数
 * @param {string}
 */
function withParser(parser) {
  if (typeof parser === 'string') {
    parser = getParser(parser);
  }

  const newCore = function(source, options) {
    if (options && !options.parser) {
      options.parser = parser;
    } else {
      options = {parser};
    }
    return core(source, options);
  };

  return enrichCore(newCore, parser);
}

/**
* The ast-types library
* @external astTypes
* @see {@link https://github.com/benjamn/ast-types}
*/
function enrichCore(core, parser) {
  // 往编译函数上挂builders和namedTypes，方便对这部分API的调用
  Object.assign(core, recast.types.namedTypes);
  Object.assign(core, recast.types.builders);
  
  // 将 Collection 中的 registerMethods 也挂到编译函数上
  core.registerMethods = Collection.registerMethods;

  core.types = recast.types;
  core.match = match;
  core.template = template(parser);

  // 将 collections 中的 filters 和 mappings 都挂在函数上
  core.filters = {};
  core.mappings = {};
  for (const name in collections) {
    if (collections[name].filters) {
      core.filters[name] = collections[name].filters;
    }
    if (collections[name].mappings) {
      core.mappings[name] = collections[name].mappings;
    }
  }
  core.use = use;
  core.withParser = withParser;
  return core;
}

module.exports = enrichCore(core, getParser());

enrichCore 就是往 core 函数上怼了一系列的 API，方便我们在写 tarnsform 中可以直接调用，不需要再去单独 import 辅助工具库。这里引出一个新的概念 Collection。

Collection

Collection 提供了一系列 jquery 式的 API 供我们操作节点。

const collections = require('./collections');

// 注册API，将 Node、JSXElement、VariableDeclarator 三种类型的 API 注册到 Collection
for (var name in collections) {
  collections[name].register();
}

VariableDeclarator 和 JSXElement 基于 Node 去扩展更多的方法。我们选择核心的 Node 来看看 register 了哪些幺蛾子，Node 的代码在 src\collections\Node.js：

'use strict';

const Collection = require('../Collection');

const matchNode = require('../matchNode');
const once = require('../utils/once');
const recast = require('recast');

const Node = recast.types.namedTypes.Node;
var types = recast.types.namedTypes;

/**
* @mixin
*/
const traversalMethods = {

  /**
   * 查找特定类型的节点
   *
   * @param {type}
   * @param {filter}
   * @return {Collection}
   */
  find: function(type, filter) {
    // ...
  },

  /**
   * 返回一个包含创建当前选定路径范围的路径的集合。对路径进行去重。
   *
   * @return {Collection}
   */
  closestScope: function() {
    // ...
  },

  /**
   * 向上遍历 AST 并找到所提供类型的最近节点
   *
   * @param {Collection}
   * @param {filter}
   * @return {Collection}
   */
  closest: function(type, filter) {
    // ...
  },

  /**
   * 查找每个选定路径的声明
   *
   * @param {function} nameGetter
   * @return {Collection}
   */
  getVariableDeclarators: function(nameGetter) {
    // ...
  },
};

function toArray(value) {
  return Array.isArray(value) ? value : [value];
}

/**
 * @mixin
 */
const mutationMethods = {
  /**
   * 替换一个节点
   *
   * @param {Node|Array<Node>|function} nodes
   * @return {Collection}
   */
  replaceWith: function(nodes) {
    // ...
  },

  /**
   * 在当前节点之前插入一个新节点
   *
   * @param {Node|Array<Node>|function} insert
   * @return {Collection}
   */
  insertBefore: function(insert) {
    // ...
  },

  /**
   * 在当前节点之后插入一个新节点。
   *
   * @param {Node|Array<Node>|function} insert
   * @return {Collection}
   */
  insertAfter: function(insert) {
    // ...
  },

  // 移除节点
  remove: function() {
    // ...
  }

};

function register() {
  // 将遍历的API注册到Collection
  Collection.registerMethods(traversalMethods, Node);
  // 将操作的API注册到Collection
  Collection.registerMethods(mutationMethods, Node);
  // 设置默认的集合类型
  Collection.setDefaultCollectionType(Node);
}

// once 跟 lodash.once 一样，重复调用同一个函数返回跟第一次调用的结果
exports.register = once(register);

这个文件定义了很多 mixin 函数，具体逻辑在下面拆解 transform 使用到了再看，先看下 Collection.registerMethods 做了什么事情？

registerMethods

registerMethods 逻辑位于 src\Collection.js：

const CPt = Collection.prototype;

/**
 * 此函数将提供的方法添加到相应类型集合的原型中。
 * 如果未传递类型，则将方法添加到 Collection.prototype 并可用于所有集合
 */
function registerMethods(methods, type) {
  // 遍历collections中的方法
  for (const methodName in methods) {
    if (!methods.hasOwnProperty(methodName)) {
      return;
    }
    // 判断指定类型中是否有方法冲突
    if (hasConflictingRegistration(methodName, type)) {
      let msg = `There is a conflicting registration for method with name "${methodName}".\nYou tried to register an additional method with `;

      if (type) {
        msg += `type "${type.toString()}".`
      } else {
        msg += 'universal type.'
      }

      msg += '\nThere are existing registrations for that method with ';

      const conflictingRegistrations = CPt[methodName].typedRegistrations;

      if (conflictingRegistrations) {
        msg += `type ${Object.keys(conflictingRegistrations).join(', ')}.`;
      } else {
        msg += 'universal type.';
      }

      throw Error(msg);
    }
    // ast-types中没有对应的类型就直接挂在 Collection.prototype 上
    if (!type) {
      CPt[methodName] = methods[methodName];
    } else {
      // 存在指定的类型，就挂到指定的类型上
      type = type.toString();
      if (!CPt.hasOwnProperty(methodName)) {
        installTypedMethod(methodName);
      }
      var registrations = CPt[methodName].typedRegistrations;
      registrations[type] = methods[methodName];
      astTypes.getSupertypeNames(type).forEach(function (name) {
        registrations[name] = false;
      });
    }
  }
}

registerMethods 把 jquery 式 API 都挂在 Collection 类原型或者类指定类型的原型上。这就能回答文章开头的第三个问题。 jscodeshift 通过在函数上挂上 registerMethods 方法，用于扩展内部 Collection 方法。

假设 🌰 中第一个函数的 console.log 需要保留，我们通过 registerMethods 注册一个 test 函数来排除第一个 path ，这时我们的 transform 代码如下：

module.exports = (fileInfo, api, options) => {
  const j = api.jscodeshift;
  
  // 将第一个 path 排除之后再去调用 remove，这里要把 collection 返回去
  function test ()  {
    this.__paths.shift()
    return this;
  }

  // 注册函数
  j.registerMethods({
    test: test
  })

  return j(fileInfo.source)
    .find(j.ExpressionStatement, {
      expression: {
        callee: {
          type: 'MemberExpression',
          object: { type: 'Identifier', name: 'console' },
        }
      },
    })
    .test()
    .remove()
    .toSource();
};

执行上述 transform 的结果如下：

export const sum = (a, b) => {
  console.log('计算下面两个数的和：', a, b);
  return a + b;
};

export const minus = (a, b) => {
  return a - b;
};

export const multiply = (a, b) => {
  return a * b;
};

export const divide = (a, b) => {
  return a / b;
};

至此，registerMethods 就基本分解完了。用一张图了解 jscodeshift 自带的 3 种 collections ，它们之间的关系如下：

有了这些方法，当执行到 transform 的 j（jscodeshift）时，相当于调用 core 函数，依次会调用 formSource -> formAST -> fromNodes -> fromPaths，这时就会创建一个 Collection 实例：

/**
 * 从节点路径数组创建一个新集合
 * @param {Array} paths 路径数组
 * @param {Collection} parent 父集合对象
 * @param {Type} type AST类型
 * @return {Collection}
 */
function fromPaths(paths, parent, type) {
  assert.ok(
    paths.every(n => n instanceof NodePath),
    'Every element in the array should be a NodePath'
  );

  return new Collection(paths, parent, type);
}

千呼万唤始出来，终看到 Collection 的真面目：

class Collection {

  /**
   * @param {Array} paths AST path数组
   * @param {Collection} parent 父Collection
   * @param {Array} types 集合中所有路径共有的类型数组。如果没有通过，它将从路径中推断出来
   * @return {Collection}
   */
  constructor(paths, parent, types) {
    assert.ok(Array.isArray(paths), 'Collection is passed an array');
    assert.ok(
      paths.every(p => p instanceof NodePath),
      'Array contains only paths'
    );
    // 父Collection
    this._parent = parent;
    // 缓存path数组
    this.__paths = paths;
    if (types && !Array.isArray(types)) {
      types = _toTypeArray(types);
      
    // 推断集合类型
    } else if (!types || Array.isArray(types) && types.length === 0) {
      types = _inferTypes(paths);
    }
    
    // Collection 的类型数组，🌰 中这里返回了 ['File', 'Node', 'Printable']
    this._types = types.length === 0 ? _defaultType : types;
  }

  // ...其他公共API这里不是重点，用上了再具体看
}

/**
 * 给定一组路径，这会推断所有路径的常见类型
 * @param {Array} paths path数组.
 * @return {Type} type AST类型
 */
function _inferTypes(paths) {
  let _types = [];

  if (paths.length > 0 && Node.check(paths[0].node)) {
    // 🌰 中 paths[0].node.type 是 File，所以节点的类型是 PredicateType
    const nodeType = types[paths[0].node.type];
    const sameType = paths.length === 1 ||
      paths.every(path => nodeType.check(path.node));

    if (sameType) {
      _types = [nodeType.toString()].concat(
        astTypes.getSupertypeNames(nodeType.toString())
      );
    } else {
      // try to find a common type
      _types = intersection(
        paths.map(path => astTypes.getSupertypeNames(path.node.type))
      );
    }
  }

  return _types;
}

Collection 构造函数就是在实例上挂了各种属性，对于没有类型的 Collection 做了一个推断，最终返回实例。🌰 中返回的实例对象如下：

生成了 collections 实例就能使用各种 API 对 nodePath 做各种操作，🌰 中分别调用了 find、remove 和 toSource。一个一个来看源码：

collection.find(j.ExpressionStatement, {
    expression: {
      callee: {
        type: 'MemberExpression',
        object: { type: 'Identifier', name: 'console' },
      }
    },
})

上一篇重构利器 jscodeshift 介绍了上述 AST 的搜索逻辑，忘记的童鞋可以回去看看！直接进到 find 函数中，看做了哪些封装：

/**
 * 在此集合的节点中查找特定类型的节点
 *
 * @param {type}
 * @param {filter}
 * @return {Collection}
 */
find: function(type, filter) {
  const paths = [];
  // 🌰 中是 visitExpressionStatement
  const visitorMethodName = 'visit' + type;

  const visitor = {};
  
  function visit(path) {
    /*jshint validthis:true */
    if (!filter || matchNode(path.value, filter)) {
      paths.push(path);
    }
    this.traverse(path);
  }

  this.__paths.forEach(function(p, i) {
    // 这里foreach传了第二个参数指定上下文，该方法是挂在 Collection 实例上。
    const self = this;
    // 定义一个访问者
    visitor[visitorMethodName] = function(path) {
      if (self.__paths[i] === path) {
        this.traverse(path);
      } else {
        return visit.call(this, path);
      }
    };
    
    // 最终还是调用 recast 上的 visit 方法
    recast.visit(p, visitor);
  }, this);

  // 返回值再去调用 fromPaths，再实例化 Collection 对象，这就是能够链式调用的原因！😏😏
  return Collection.fromPaths(paths, this, type);
}

🌰 中对应的 paths、this、type 如下截图所示：可以看到，find 根据参数定义一个访问者，然后调用 recast.visit。最后递归遍历找出满足 filter 条件的节点，最后再调用 Collection.fromPaths 再生成 Collection 实例。这就是第一个问题的答案——因为每个函数都返回了 Collection 实例，当然都能够调用 Collection 的方法。不同的是，此时 parent 有值并且 types 可以从 paths 中获取到。
接下来是 remove 方法：

// src\collections\Node.js
remove: function() {
  // path.prune 是 ast-types 里的函数，用于删除path
  return this.forEach(path => path.prune());
}

// src\Collection.js
class Collection {
  // ...
  forEach(callback) {
    this.__paths.forEach(
      (path, i, paths) => callback.call(path, path, i, paths)
    );
    return this;
  }
  // ...
}

remove 就比较简单了，通过遍历 paths 调用 prune（ast-types 方法）来删除 nodePath，再返回实例。最后一个方法 toSource：

// src\Collection.js
class Collection {
  // ...
  toSource(options) {
    // 找到根节点
    if (this._parent) {
      return this._parent.toSource(options);
    }
    if (this.__paths.length === 1) {
      return recast.print(this.__paths[0], options).code;
    } else {
      return this.__paths.map(p => recast.print(p, options).code);
    }
  }
  // ...
}

toSource 先是找到根节点，然后调用 recast.print 打印转换之后的代码。 options 参数可以参考 recast options 🧐🧐

testUtils

testUtils 作用就是通过模板简化用例写法。忘记用例写法的童鞋回去看下重构利器 jscodeshift ！ 😎😎 这部分代码比较简单，直接上 testUtils 流程图：

简述上图流程：

入口通过 runTest 定义用例名称，这个名称作为 describe 和 test 的名称；
执行 runTest，通过参数获取 parser，然后读入测试输入文件（remove_console.input.js）、输出文件（remove_console.output.js）内容，并 require transform 模块；
调用 runInlineTest：获取 transform 代码和 jscodeshift 函数，如果有指定 parser，就调用 withParser 生成解析器；最后调用 transform 输出转换结果；
将结果与 output 文件内容进行对比，输出测试结果。

总结

最后结合下图做一个总结： jscodeshift 层次还是很清晰的，最底层的编译器 recast 和 AST 的操作工具 ast-types，核心逻辑 core 搭载很多的工具方法（模板、解析器的多样化处理、集合概念、节点匹配工具），应用层有 CLI 和测试工具。文章从 transform 作为入口而没有从命令行入口开始解析，是因为命令行应用的逻辑都是基础操作，而核心部分在 core 的 Collection 和 parser！
我之前参加公司的代码训练时，讲师有分享通过 registXXX 去扩展核心功能从而满足 SOLID 原则。那会觉得扯淡，实际编码很难想到吧！！对不起，是我格局小了，这也是看源码的好处之一吧，模仿着、模仿着，你就会了！你会了吗？