关于如何实现一个 Babel 的自定义插件之自动埋点

「这是我参与2022首次更文挑战的第16天，活动详情查看：2022首次更文挑战」

PS：最近在看光神的 Babel 插件通关秘籍，想把自己学习的一些过程记录下来。

@babel/core

@babel/core 包是完成整个编译流程，从源码到目标代码，生成 sourcemap。

transformFromAstSync(
  ast,
  sourceCode,
  options
); // => { code, map, ast }

@babel/parser

对源码进行 parse，可以通过 plugins、sourceType 等来指定 parse 语法。

sourceType： 指定是否支持解析模块语法，有 module、script、unambiguous 3个取值，module 是解析 es module 语法，script 则不解析 es module 语法，当作脚本执行，unambiguous 则是根据内容是否有 import 和 export 来确定是否解析 es module 语法。

---

示例代码 sourceCode.js

import example1 from 'example1';
import * as example2 from 'example2';
import { example3 } from 'example3';
import 'example4';

function example5() {
  console.log('example5');
}

class Example6 {
  example2() {
    return 'example6';
  }
}

const example7 = () => 'example7';

const example8 = function () {
  console.log('example8');
}

入口文件，index.js。

通过 fs 读取文件内容，通过 @babel/parser 中 parser 解析 code，变为 AST，transformFromAstSync 函数配置 plugin。

const { transformFromAstSync } = require('@babel/core');
const parser = require('@babel/parser');
const autoTrackPlugin = require('./plugin/auto-track-plugin');
const fs = require('fs');
const path = require('path');

const sourceCode = fs.readFileSync(path.join(__dirname, './sourceCode.js'), {
  encoding: 'utf-8'
});

const ast = parser.parse(sourceCode, {
  sourceType: 'unambiguous'
});

const { code } = transformFromAstSync(ast, sourceCode, {
  plugins: [[autoTr ackPlugin, {
    trackerPath: 'tracker'
  }]]
});

console.logs(code);

搭建基本框架 auto-track-plugin.js

visitor模式

当被操作的对象结构比较稳定，而操作对象的逻辑经常变化的时候，通过分离逻辑和对象结构，使得他们能独立扩展。这就是 visitor 模式的思想。对应到 babel traverse 的实现，就是 AST 和 visitor 分离，在 traverse（遍历）AST 的时候，调用注册的 visitor 来对其进行处理。

const { declare } = require('@babel/helper-plugin-utils');

const autoTrackPlugin = declare((api, options) => {
  api.assertVersion(7);

  return {
    visitor: {
      //....
    }
  }
});

module.exports = autoTrackPlugin;

模块引入

在 Program 根结点里，通过 enter 的时候判断 state 上是否存在 trackerImportId，即是否已经导入，没有就引入 tracker 模块，用 generateUid 生成唯一 id，然后放到 state。

visitor: {
  Program: {
    enter(path, state) {
      if (!state.trackerImportId) {
        state.trackerImportId = importModule.addDefault(path, 'tracker', {
          nameHint: path.scope.generateUid('tracker')
        }).name
      }
    }
  },
}

埋点代码导入之后，下一步就是向函数体中插入代码。

找到对应的函数，这里要处理的有：ClassMethod、ArrowFunctionExpression、FunctionExpression、FunctionDeclaration 这些节点。

visitor: {
  Program: {
    enter(path, state) {
      if (!state.trackerImportId) {
        state.trackerImportId = importModule.addDefault(path, 'tracker', {
          nameHint: path.scope.generateUid('tracker')
        }).name;
      }
    }
  },
  'ClassMethod|ArrowFunctionExpression|FunctionExpression|FunctionDeclaration'(path, state) {
    const bodyPath = path.get('body');
    if (bodyPath.isBlockStatement()) {
      console.log('插入埋点代码');
    }
  }
}

我们发现，正好是执行了三次，对应 source 中三个有函数体的代码。如果不存在函数体，包装一下，然后修改下 return 值。

'ClassMethod|ArrowFunctionExpression|FunctionExpression|FunctionDeclaration'(path, state) {
  const bodyPath = path.get('body');
  if (bodyPath.isBlockStatement()) {
    console.log('插入埋点代码');
  } else {
    const ast = api.template.statement(`{return PREV_BODY;}`)({ PREV_BODY: bodyPath.node });
    bodyPath.replaceWith(ast);
    console.log('插入埋点代码');
  }
}

执行了四次，没有函数体的代码也被替换成功了。

visitor: {
  Program: {
    enter(path, state) {
      if (!state.trackerImportId) {
        state.trackerImportId = importModule.addDefault(path, 'tracker', {
          nameHint: path.scope.generateUid('tracker')
        }).name;
        state.trackerAST = api.template.statement(`${state.trackerImportId}()`)();
      }
    }
  },
  'ClassMethod|ArrowFunctionExpression|FunctionExpression|FunctionDeclaration'(path, state) {
    const bodyPath = path.get('body');
    if (bodyPath.isBlockStatement()) {
      bodyPath.node.body.unshift(state.trackerAST)
    } else {
      const ast = api.template.statement(`{${state.trackerImportId}();return PREV_BODY;}`)({ PREV_BODY: bodyPath.node });
      bodyPath.replaceWith(ast);
    }
  }
}

完整代码：

const { declare } = require('@babel/helper-plugin-utils');
const importModule = require('@babel/helper-module-imports');

const autoTrackPlugin = declare((api, options, dirname) => {
  api.assertVersion(7);

  return {
    visitor: {
      Program: {
        enter(path, state) {
          if (!state.trackerImportId) {
            state.trackerImportId = importModule.addDefault(path, 'tracker', {
              nameHint: path.scope.generateUid('tracker')
            }).name;
            state.trackerAST = api.template.statement(`${state.trackerImportId}()`)();
          }
        }
      },
      'ClassMethod|ArrowFunctionExpression|FunctionExpression|FunctionDeclaration'(path, state) {
        const bodyPath = path.get('body');
        if (bodyPath.isBlockStatement()) {
          bodyPath.node.body.unshift(state.trackerAST);
        } else {
          const ast = api.template.statement(`{${state.trackerImportId}();return PREV_BODY;}`)({ PREV_BODY: bodyPath.node });
          bodyPath.replaceWith(ast);
        }
      }
    }
  }
});
module.exports = autoTrackPlugin;

将埋点之后的代码，输出到另一个文件夹中。

index.js

const { transformFromAstSync } = require('@babel/core');
const parser = require('@babel/parser');
const autoTrackPlugin = require('./plugin/auto-track-plugin');
const fs = require('fs');
const fse = require('fs-extra');
const path = require('path');

const sourceCode = fs.readFileSync(path.join(__dirname, './sourceCode.js'), {
  encoding: 'utf-8'
});

const ast = parser.parse(sourceCode, {
  sourceType: 'unambiguous'
});

const outputDir = path.resolve(__dirname, './lib');

const { code } = transformFromAstSync(ast, sourceCode, {
  plugins: [[autoTrackPlugin, {
    trackerPath: 'tracker'
  }]]
});

fse.ensureDirSync(outputDir);
fse.writeFileSync(path.join(outputDir, 'lib.js'), code);

lib/index.js 生成埋点之后的文件。