使用 babel ast 做代码转换

前言

现在对 ast(abstract syntax tree 虚拟语法树)应该也不陌生了。像 vue react 开发所写的非浏览器标准的代码,最后都是转化成符合标准的代码,都是中间靠这个处理。也包括小程序,react-native 这些也是要经过同样处理转换。

ast 三板斧:parse(解析)-> transform(转换)-> generate(生成)。
听起来有点像变形金刚,我在想法术里的变身是不是原理也是这样:实体拆散成原子,原子转换变形,然后再重新组合。哈哈。

ast 工具其实很多,babel 应该是前端接触较多较早的一个,毕竟很长一段时间内都是它在帮着处理前端的编译。babel 对于上面三板斧拆成了三个工具:@babel/parser@babel/traverse@babel/generator
乍看之下,其它两个都对得上,好像 transform 不太对。其实转换规则是自己定的,babel 也不知道使用者要转成啥,所以它提供了一个 traverse 工具,用来方便获取节点。因为 ast 是一串层级深又复杂的对象,很难自己手动找到想要的节点,所以提供了一个遍历工具。

当然有个遍历工具还是不够,最好还是搭配一个很牛网站:https://www.astexplorer.net 可以直接左边输入源代码,右边看结构树。从这个网站里可以看到其实还有很多别的 ast 工具,每个工具都有自己的节点定义。

实践

现在尝试进行一段代码的转换:

// input:
import Input from './Input';
export const Component = (props) => <div><Input onChange={props.onChange} /></div>;

// output:
import { Input } from './Input';
const map = {
  run: () => {}
};
export default function Component(props) {
  return <div><Input onChange={map[props.method]} /></div>;
}

就是,1、把 import default 导出变成解构导出,2、把 export 分散导出变成 default 导出,3、把事件的触发函数变成调用一个事件映射的属性。

过程

代码在最下面

解析和生成都没啥说的,主要还是这个转换的实现。

1、寻找节点

看一下遍历的代码:

traverse(ast, {
  enter(path) {
    if (path.isIdentifier({ name: "n" })) {
      path.node.name = "x";
    }
  },
  FunctionDeclaration: function(path) {
    path.node.id.name = "x";
  },
});

可以在 enter 里加上一些条件判断,找到相应的节点。但是这样只能找到一些普通的节点,定向的还是得下面的 FunctionDeclaration 这种类型找法。

这时候就可以先把源代码在 astexplorer 上解析一遍。上面有 autofocus 功能,鼠标在左侧的代码节点移动,右边可以定向到相应的节点结构。这时候节点的 type 就是相应的类型。

比如在这个实践案例中,我们要找的是箭头函数返回的那个 jsx,那么它的 type 就是 ArrowFunctionExpression。当然这是假设只有一个箭头函数,实际应用起来要加别的判断来确定是要找的那个。

当然 path 不只是 isIdentifier 一个方法,本示例也可以写成 path.isArrowFunctionExpression。基本上所有节点类型都有相应的 isType 的方法,可以点进去它的类型声明,看看有哪些方法。

2、拼接节点

代码中需要生成一些补偿代码或者修改一些节点内容,比如实践示例:import export 的导出变化,依然借助 astexplorer。把心目中的转换后的代码放上去解析,看看要转成这样,需要加什么节点。比如示例里的箭头函数变成 function 函数,就要补充 { return } 代码。
这在 ast 中就是包裹一个 BlockStatement 和 ReturnStatement 对象,当然实际应用中要判断这个箭头函数是直接返回还是 return 返回。至于这些节点对象都有些什么参数需要怎么用,就直接看解析后本该是咋样,给相应填充上去就行。简单来讲,就是拼接出一个 ast 节点。

3、生成节点

有时需要生成一大段辅助函数插入进去,这时候用拼接 ast 节点的方法就太麻烦了,可能是个非常复杂的对象。在知道这段代码或者说能单独拼出一段代码的前提下,可以把这段代码用 parser 的方法,解析成一个节点 ast,加入到原来的主体 ast 里。比如实践示例里:要加入一个 map 对象,就可以将其转成 ast 后,和之前箭头函数组合成一个 function 函数。

调试

因为是 node 工具,还是建议装个 node inspector 工具,借助 chrome 的 CDP 协议把信息打印到浏览器来,可以看到完整的对象。不然看终端打印的信息太麻烦了。

示例代码

const parser = require('@babel/parser');
const traverse = require('@babel/traverse');
const generate = require('@babel/generator');

const source = `
import Input from './Input';
export const Component = (props) => <div><Input onChange={props.onChange} /></div>;
`;

(async () => {
  const ast = parser.parse(source, {
    sourceType: 'module',
    plugins: [
      'jsx',
      'flow', // 官网示例,需要支持 ts,可以改成 'typescript',但与 'flow' 冲突。
    ],
  });

  let jsxElementPath = null;
  let onChangePath = null;
  traverse.default(ast, {
    enter(path) {
      if (path.isIdentifier({ name: 'onChange' })) {
        path.node.name = 'method';
        onChangePath = path;
      }
    },
    ArrowFunctionExpression: function (path) {
      jsxElementPath = path;
    },
  });

  // props.onChange 这个节点
  const propsOnChange = onChangePath.parent;
  // onChange={props.onChange} 这个节点
  const memberExpression = onChangePath.parentPath.parent.expression;

  // computed 要加,不然会报错
  memberExpression.computed = true;
  memberExpression.property = {
    computed: false,
    type: 'MemberExpression',
    object: {
      type: 'Identifier',
      name: propsOnChange.object.name,
    },
    property: {
      type: 'Identifier',
      name: propsOnChange.property.name,
    },
  };
  // 上面对象是引用的,把改名字放到最后,否则影响上面的名字
  memberExpression.object.name = 'map';

  // 示例获取最上面的 import 语句 ast,实际应用要加别的判断辅助
  const importDeclaration = ast.program.body[0];
  importDeclaration.specifiers[0] = {
    type: 'ImportSpecifier',
    imported: {
      type: 'Identifier',
      name: 'Input'
    },
    local: {
      type: 'Identifier',
      name: 'Input'
    },
  };

  // 增加代码以 ast 加入的方式,不要用插入文本
  const functionAst = parser.parse(`
  const map = {
    run: () => {},
  };
  `, {
    sourceType: 'module',
    plugins: [
      'jsx',
      'flow',
    ],
  });

  // 简单查找相应的导出语句,真实情况要附加多点判断去寻找
  const exportDeclaration = ast.program.body.pop();
  const componentName = exportDeclaration.declaration.declarations[0].id.name;

  // 生成 export default 导出的 ast 节点
  const exportDefaultDeclaration = {
    type: 'ExportDefaultDeclaration',
    declaration: {
      type: 'FunctionDeclaration',
      id: {
        type: 'Identifier',
        name: componentName,
      },
      params: [...jsxElementPath.node.params],
      body: {
        type: 'BlockStatement',
        body: [{
          type: 'ReturnStatement',
          argument: jsxElementPath.node.body,
        }],
      },
    },
  }

  const { code, map } = generate.default(
    {
      type: 'Program',
      body: [].concat(importDeclaration, functionAst, exportDefaultDeclaration),
    },
    { sourceMaps: true },
  );

  console.info('// input:');
  console.info(source);
  console.info('// output:\n',);
  console.info(code);
})();