手把手教写 TypeScript Transformer Plugin

手把手教写 TypeScript Transformer Plugin

用过 ant-design 的同学可能对 babel-plugin-import 有印象,它可以帮助实现模块的按需引用,比如:

import { Button } from 'antd'

在使用该 Plugin 之后会被转换成:

import Button from 'antd/lib/button'

在一个没有使用 antd 全部组件的项目里,这样做可以明显减少打包后的代码体积。
可是,如果你在一个没有使用 Babel 的 TypeScript 项目里,想要实现类似的功能,该怎么办呢?

这就要用到本文的主角:custom transformation,这是从 TypeScript@2.3 开始引入的新能力,他让我们可以部分修改 TS 从源码转换成的语法树,从而控制生成的 JavaScript 代码,最终完成上述的转换。

先让我们从 TS 中代码语法树的样子说起。

预备知识

1. 抽象语法树(AST)

AST 是为了方便计算机理解源代码、用于表达源代码语法结构的树状结构,由称作节点(Node)的数据结构组成。

例如:

const name: string = 'Tom'

上面这段代码在 TS 会解析成下图所示的 AST:

确切来说,上图实际上是语法树而不是抽象语法树,因为节点里面仍然包含了「冒号」等多余信息,还不够「抽象」,但是,因为在之后处理的过程中实际面对的就是这样的语法树,因此在这里不做严格的区分。

TS 中所有 AST 的根节点都是 SourceFile,顾名思义,这是一个附加了源文件信息的 AST 节点(Node)。

源码中只有一个变量声明语句,该声明生成了以下结构:

  • 表示这是一个常量声明的 ConstKeyword 节点
  • 表达变量名的 Identifier 节点
  • 表达变量类型的 StringKeyword 节点
  • 表达变量初始值的 StringLiteral 节点
  • 其他附属信息节点

TypeScript/typescript.d.ts 源码中,用枚举类型 SyntaxKind 定义了所有的 AST 节点类型,到目前为止近 300 个,可以看出来 AST 的树形结构非常得精确细致,想手动分析记忆比较困难,可以借助 AST explorer 这个可视化工具帮助理解代码的 AST 结构。

2. TS 编译流程

和 Babel 以及其他编译到 JavaScript 的工具类似,TS 的编译流程包含以下三步:

解析->转换->生成

包含了以下几个关键部分:

  • Scanner:从源码生成 Token
  • Parser:从 Token 生成 AST
  • Binder:从 AST 生成 Symbol
  • Checker:类型检查
  • Emitter:生成最终的 JS 文件

图示如下:

我们的标题中所指的 transformer Plugin 就是在 Emitter 阶段起作用。

transformer Plugin 如何启用?

tsc命令不支持直接配置 transformer 的参数,你可以手动引入 typescript 来自己编译,当然,目前最方便的办法是在 Webpack + ts-loader 的项目中,给 ts-loader 配置 getCustomTransformers 选项:

{
  test: /\.tsx?$/,
  loader: 'ts-loader',
  options: {
    ... // other loader's options
    getCustomTransformers: () => ({ before: [yourImportedTransformer] })
  }
}

详见 ts-loader 文档。

实际编写一个 transformer Plugin

目标

我们的目标就是实现文章开头代码示例中的转换:

// before
import { Button } from 'antd'

// after
import Button from 'antd/lib/button'

了解需要改什么

Custom Transformer 操作是 AST,所以我们需要了解代码转换前后的 AST 区别在哪里。

转换前:

import { Button } from 'antd'

代码的 AST 如下:

转换后:

import Button from 'antd/lib/button'

代码的 AST 如下:

可以看出,我们需要做的转换有两处:

  • 替换 ImportClause 的子节点,但保留其中的 Identifier
  • 替换 StringLiteral 为原来的值加上上面的 Identifier

那么,该如何找到并替换对应的节点呢?

如何遍历并替换节点

TS 提供了两个方法遍历 AST:

  • ts.forEachChild
  • ts.visitEachChild

两个方法的区别是:

forEachChild 只能遍历 AST,visitEachChild 在遍历的同时,提供给此方法的 visitor 回调的返回节点,会被用来替换当前遍历的节点,因此我们可以利用 visitEachChild 来遍历并替换节点。

先看一下这个方法的签名:

/**
 * Visits each child of a Node using the supplied visitor, possibly returning a new Node of the same kind in its place.
 *
 * @param node The Node whose children will be visited.
 * @param visitor The callback used to visit each child.
 * @param context A lexical environment context for the visitor.
 */
function visitEachChild<T extends Node>(node: T, visitor: Visitor, context: TransformationContext): T

假设我们已经拿到了 AST 的根节点 SourceFile 和 TransformationContext,我们就可以用以下代码遍历 AST:

ts.visitEachChild(SourceFile, visitor, ctx)

function visitor(node) {
  if(node.getChildCount()) {
    return ts.visitEachChild(node, visitor, ctx)
  }
  return node
}

注意:visitor 的返回节点会被用来替换 visitor 正在访问的节点。

如何创建节点

TS 中 AST 节点的工厂函数全都以 create 开头,在编辑器里敲下:ts.create,代码补全列表里就能看到很多很多和节点创建有关的方法:

比如,创建一个 1+2 的节点:

ts.createAdd(ts.createNumericLiteral('1'), ts.createNumericLiteral('2'))

如何判断节点类型

前面说过,ts.SyntaxKind里存储了所有的节点类型。同时,每个节点中都有一个 kind 字段标明它的类型。我们可以用以下代码判断节点类型:

if(node.kind === ts.SyntaxKind.ImportDeclaration) {
  // Get it!
}

也可以用 ts-is-kind 模块简化判断:

import * as kind from 'ts-is-kind'
if(kind.isImportDeclaration(node)) {
  // Get it!
}

那么,我们之前的 visitor 就可以继续补充下去:

import * as kind from 'ts-is-kind'
function visitor(node) {
  if(kind.isImportDeclaration(node)) {
    const updatedNode = updateImportNode(node, ctx)
    return updateNode
  }
  return node
}

因为 Import 语句不能嵌套在其他语句下面,所以 ImportDeclaration 只会出现在 SourceFile 的下一级子节点上,因此上面的代码并没有对 node 做深层递归遍历。

只要 updateImportNode 函数完成了之前图中表现出的 AST 转换,我们的工作就完成了。

如何更新 ImportDeclaration 节点

下面关注 updateImportNode 怎么实现。

我们已经拿到了 ImportDeclaration 节点,还记得到底要干什么吗?

  • 用 Identifier 替换 NamedImports 的子节点
  • 修改 StringLiteral 的值

为了方便找到需要的节点,我们对 ImportDeclaration 做递归遍历,只对 NamedImports 和 StringLiteral 做特殊处理:

function updateImportNode(node: ts.Node, ctx: ts.TransformationContext) {
  const visitor: ts.Visitor = node => {
    if (kind.isNamedImports(node)) {
      // ...
    }

    if (kind.isStringLiteral(node)) {
      // ...
    }

    if (node.getChildCount()) {
      return ts.visitEachChild(node, visitor, ctx)
    }
    return node
  }
}

首先处理 NamedImports。

在 AST explorer 的帮助下,可以发现 NamedImports 包含了三部分,两个大括号和一个叫 Button 的 Identifier,我们在 isNamedImports 的判断下,直接返回这个 Identifier,就可以取代原先的 NamedImports:

if (kind.isNamedImports(node)) {
   const identifierName = node.getChildAt(1).getText()
  // 返回的节点会被用于取代原节点
  return ts.createIdentifier(identifierName)
}

再处理 StringLiteral。

发现要返回新的 StringLiteral,要用到 isNamedImports 判断里提取出来的 identifierName。因此我们先把 identifierName 提取到外层定义,作为 updateImportNode 的内部状态。

同时,antd/lib 目录下的文件名没有大写字母,因此要把 identifierName 中首字母大写去掉:

if (kind.isStringLiteral(node)) {
  const libName = node.getText().replace(/[\"\']/g, '')
  if (identifierName) {
    const fileName = camel2Dash(identifierName)
    return ts.createLiteral(`${libName}/lib/${fileName}`)
  }
}

// from: https://github.com/ant-design/babel-plugin-import
function camel2Dash(_str: string) {
  const str = _str[0].toLowerCase() + _str.substr(1)
  return str.replace(/([A-Z])/g, ($1) => `-${$1.toLowerCase()}`)
}

完整的 updateImportNode 实现如下:

function updateImportNode(node: ts.Node, ctx: ts.TransformationContext) {
  const visitor: ts.Visitor = node => {
  if (kind.isNamedImports(node)) {
    const identifierName = node.getChildAt(1).getText()
      return ts.createIdentifier(identifierName)
  }

    if (kind.isStringLiteral(node)) {
    const libName = node.getText().replace(/[\"\']/g, '')
    if (identifierName) {
      const fileName = camel2Dash(identifierName)
      return ts.createLiteral(`${libName}/lib/${fileName}`)
    }
  }

    if (node.getChildCount()) {
      return ts.visitEachChild(node, visitor, ctx)
    }
    return node
  }
}

以上,我们就成功实现了如下代码转换:

// before
import { Button } from 'antd'

// after
import Button from 'antd/lib/button'

以上代码整合起来,就是一个完整的 Transformer Plugin,完整代码请见:newraina/learning-ts-transfomer-plugin

改进

刚才实现的只是一个最最精简的版本,距离 babel-plugin-import 的完整功能还有很远,比如:

  • 同时 Import 多个组件怎么办,如import { Button, Alert } from 'antd'
  • Import 时用 as 重命名了怎么办,如import { Button as Btn } from 'antd'
  • 如果 CSS 也要按需引入怎么办

以上都可以在 AST explorer 的帮助下找到 AST 转换前后的区别,然后按照本文介绍的流程实现。

附注