关于如何实现一个 Babel 的自定义插件之自动生成 API 文档

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

使用 vsCode 编辑器，打入 /** 时，按下回车

编辑器会根据函数的入参和出参情况，自动生成注释。

sourceCode.ts

/**
 * example1 测试代码
 * @param name 名字
 * @param age 年龄
 * @param sex 性别
 * @returns string
 */
function example1(name: string, age: number, sex: boolean): string {
  return `hi, ${name}, ${age}, ${sex}`;
}

/**
 * 类测试
 */
class Guang {
  name: string; // name 属性
  constructor(name: string) {
    this.name = name;
  }

  /**
   * 方法测试
   */
  sayHi(): string {
    return `hi, I'm ${this.name}`;
  }
}

使用 fs.readFileSync 读取文件内容，通过 babel parser 把文件内容 parse 成 ast，使用 transformFromAstSync 转化 ast。

index.js

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

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

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

const { code } = transformFromAstSync(ast, sourceCode, {
  plugins: [[autoDocumentPlugin, {
    outputDir: path.resolve(__dirname, './docs'),
    format: 'json'
  }]]
});

console.log(code);

插件的基本结构:

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

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

  return {
    pre(file) { },
    visitor: {
      FunctionDeclaration(path, state) { },
      ClassDeclaration(path, state) { }
    },
    post(file) { }
  }
});

module.exports = autoDocumentPlugin;

在全局的 file 对象中放一个 docs 的数组，用于收集信息。

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

  return {
    pre(file) {
      file.set('docs', [])
    },
    visitor: {
      FunctionDeclaration(path, state) { },
      ClassDeclaration(path, state) { }
    },
    post(file) {
      const docs = file.get('docs');
    }
  }
});

需要处理 FunctionDeclaration 节点和 ClassDelcaration 节点

因为 docs 数组是我们自行创建的，所以我们需要什么样数据，就往数组里面 push 什么样数据。

FunctionDeclaration 有两个参数，path 和 state。在 path 中，可以获取到很多我们需要的东西。

对于一个函数，里面的实现细节我们是不关心的。外部调用函数，大体上只需要知道该函数的函数名，参数，返回值类型，说明信息差不多就是这些。

按照这个思路，进行代码编写。

FunctionDeclaration(path, state) {
  const docs = state.file.get('docs');
  docs.push({
    type: 'function',
    name: path.get('id'),
    params: path.get('params').map(paramPath => {
      return {
        name: path.get('id'),
        type: path.get('returnType')
      }
    }),
    return: path.get('returnType'),
    doc: path.node.leadingComments
  })
},

将获取到的 docs 打印出来，发现这个 docs 压根不能使用，需要继续对数据进行处理。

FunctionDeclaration(path, state) {
  const docs = state.file.get('docs');
  docs.push({
    type: 'function',
    name: path.get('id').toString(),
    params: path.get('params').map(paramPath => {
      return {
        name: path.get('id').toString(),
        type: path.get('returnType').getTypeAnnotation()
      }
    }),
    return: path.get('returnType').getTypeAnnotation(),
    doc: path.node.leadingComments[0].value
  })
},

现在似乎好像还可以，大致算是取到了我们需要的数据

FunctionDeclaration(path, state) {
  const docs = state.file.get('docs');
  docs.push({
    type: 'function',
    name: path.get('id').toString(),
    params: path.get('params').map(paramPath => {
      return {
        name: paramPath.toString(),
        type: resolveType(paramPath.getTypeAnnotation())
      }
    }),
    return: returnType(path.get('returnType').getTypeAnnotation()),
    doc: path.node.leadingComments && parseComment(path.node.leadingComments[0].value)
  });
  state.file.set('docs', docs);
},

• 注释信息用 doctrine 来 parse

function resolveType(tsType) {
  const typeAnnotation = tsType.typeAnnotation;
  if (!typeAnnotation) {
    return;
  }

  switch (typeAnnotation.type) {
    case 'TSStringKeyword':
      return 'string';
    case 'TSNumberKeyword':
      return 'number';
    case 'TSBooleanKeyword':
      return 'boolean';
  }
}

function returnType(tsType) {
  switch (tsType.type) {
    case 'TSStringKeyword':
      return 'string';
    case 'TSNumberKeyword':
      return 'number';
    case 'TSBooleanKeyword':
      return 'boolean';
  }
}

function parseComment(commentStr) {
  if (!commentStr) {
    return;
  }

  return doctrine.parse(commentStr, {
    unwrap: true
  });
}

ClassDeclaration 的处理

提取 constructor、method、properties 的信息。

ClassDeclaration(path, state) {
  const docs = state.file.get('docs');
  const classInfo = {
    type: 'class',
    name: path.get('id').toString(),
    constructorInfo: {},
    methodsInfo: [],
    propertiesInfo: [],
    doc: path.node.leadingComments && parseComment(path.node.leadingComments[0].value)
  };
  path.traverse({
    ClassProperty(path) {
      classInfo.propertiesInfo.push({
        name: path.get('key').toString(),
        type: resolveType(path.getTypeAnnotation()),
        doc: [path.node.leadingComments, path.node.trailingComments].filter(Boolean).map(comment => {
          return parseComment(comment.value);
        }).filter(Boolean)
      })
    },
    ClassMethod(path) {
      if (path.node.kind === 'constructor') {
        classInfo.constructorInfo = {
          params: path.get('params').map(paramPath => {
            return {
              name: paramPath.toString(),
              type: resolveType(paramPath.getTypeAnnotation()),
              doc: parseComment(path.node.leadingComments[0].value)
            }
          })
        }
      } else {
        classInfo.methodsInfo.push({
          name: path.get('key').toString(),
          doc: parseComment(path.node.leadingComments[0].value),
          params: path.get('params').map(paramPath => {
            return {
              name: paramPath.toString(),
              type: resolveType(paramPath.getTypeAnnotation())
            }
          }),
          return: resolveType(path.getTypeAnnotation())
        })
      }
    }
  });
  docs.push(classInfo);
  state.file.set('docs', docs);
}

在 post 阶段就能拿到所有的信息了，之后就是文档的生成。

JSON

比如我们希望生成 JSON 格式。在 post 阶段拿到传入的参数，使用 generate 函数生成结果，调用 fs.writeFileSync 写入提前设置好的路径。

post(file) {
  const docs = file.get('docs');
  const res = generate(docs, options.format);
  fse.ensureDirSync(options.outputDir);
  fse.writeFileSync(path.join(options.outputDir, 'docs' + res.ext), res.content);
}

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

function generate(docs, format = 'json') {
  if (format === 'json') {
    return {
      ext: '.json',
      content: renderer.json(docs)
    }
  }
}

renderer/json.js

module.exports = function (docs) {
  return JSON.stringify(docs, null, 2);
}

docs/docs.json

[
  {
    "type": "function",
    "name": "example1",
    "params": [
      {
        "name": "name",
        "type": "string"
      },
      {
        "name": "age",
        "type": "number"
      },
      {
        "name": "sex",
        "type": "boolean"
      }
    ],
    "return": "string",
    "doc": {
      "description": "example1 测试代码",
      "tags": [
        {
          "title": "param",
          "description": "名字",
          "type": null,
          "name": "name"
        },
        {
          "title": "param",
          "description": "年龄",
          "type": null,
          "name": "age"
        },
        {
          "title": "param",
          "description": "性别",
          "type": null,
          "name": "sex"
        }
      ]
    }
  },
  {
    "type": "class",
    "name": "Guang",
    "constructorInfo": {
      "params": [
        {
          "name": "name",
          "type": "string",
          "doc": {
            "description": "name 属性",
            "tags": []
          }
        }
      ]
    },
    "methodsInfo": [
      {
        "name": "sayHi",
        "doc": {
          "description": "方法测试",
          "tags": []
        },
        "params": []
      }
    ],
    "propertiesInfo": [
      {
        "name": "name",
        "type": "string",
        "doc": []
      }
    ],
    "doc": {
      "description": "类测试",
      "tags": []
    }
  }
]

Markdown

renderer/markdown.js

根据 type 的不同，进行不同 markown 的拼接

module.exports = function (docs) {
  let str = '';

  docs.forEach(doc => {
    if (doc.type === 'function') {
      str += '## ' + doc.name + '\n\n';
      str += doc.doc.description + '\n\n';
      if (doc.doc.tags) {
        doc.doc.tags.forEach(tag => {
          str += tag.name + ': ' + tag.description + '\n\n';
        })
      }
      str += '> ' + doc.name + '(';
      if (doc.params) {
        str += doc.params.map(param => {
          return param.name + ': ' + param.type;
        }).join(', ');
      }
      str += ')\n';
      str += '#### 参数: \n\n';
      if (doc.params) {
        str += doc.params.map(param => {
          return '- ' + param.name + '(' + param.type + ')';
        }).join('\n');
      }
      str += '\n\n'
    } else if (doc.type === 'class') {
      str += '## ' + doc.name + '\n\n';
      str += doc.doc.description + '\n';
      if (doc.doc.tags) {
        doc.doc.tags.forEach(tag => {
          str += tag.name + ': ' + tag.description + '\n';
        })
      }
      str += '> new ' + doc.name + '(';
      if (doc.params) {
        str += doc.params.map(param => {
          return param.name + ': ' + param.type;
        }).join(', ');
      }
      str += ')\n';
      str += '#### 属性: \n\n';
      if (doc.propertiesInfo) {
        doc.propertiesInfo.forEach(param => {
          str += '- ' + param.name + ': ' + param.type + '\n';
        });
      }
      str += '#### 方法: \n\n';
      if (doc.methodsInfo) {
        doc.methodsInfo.forEach(param => {
          str += '- ' + param.name + '\n';
        });
      }
      str += '\n'
    }
    str += '\n'
  })
  return str;
}

传入的 format 为 markdown 时

transformFromAstSync(ast, sourceCode, {
  plugins: [[autoDocumentPlugin, {
    outputDir: path.resolve(__dirname, './docs'),
    format: 'markdown'
  }]]
});

改写 generate 函数

function generate(docs, format = 'json') {
  if (format === 'markdown') {
    return {
      ext: '.md',
      content: renderer.markdown(docs)
    }
  } else {
    return {
      ext: '.json',
      content: renderer.json(docs)
    }
  }
}

---

大功告成了 🎉🎉🎉🎉🎉🎉🎉