asterisc

技術ネタ中心にその他雑多なこと

TypeScriptのcompiler APIをいじる

はじめに

これを作っているときに、もともとのagreedとswaggerの型の帳尻を合わせるために、TypeScriptのASTをいじっていろいろやりました。 そのときに調べたことなどをメモ代わりにまとめます。この内容は調べたこと+僕独自の解釈が入っている可能性があり、正確ではないかもしれませんが、ご了承ください。

TypeScript Compiler API

何ができるようになるか

TypeScript Compilerを使うことにより、TypeScriptのASTの解析や編集などを行うことができます。linterや、静的解析、自動生成系のツールを作る際によく使う機能です。 例えば前述のagreed-typedではTypeScriptの型定義やコメントをASTを解析することにより抜き出し、そこからOpenAPI2.0(Swagger)を生成しています。 実際のアプリケーションコードを書く際には全く触れることはない(また、触れるべきではありません)機能ですが、用途を絞って活用することにより、利便性の高いツール群を実装することができます。

サンプル実装

TypeScriptのCompilerAPIを使って、アプリケーションコードの中からassertを削除するツールを書いてみようと思います。 以降unassertと呼びますが、unassertやassertionについては過去の記事を参照していただくとわかりやすいと思います。

akito0107.hatenablog.com

実装にあたって、ここに公式のdocumentがありますが、そこまで詳細なdocumentationではなく、ここにある情報だけで何かを実装するのは少し難しいかなと思っています。このAPIをいじる際には、GithubのTypeScriptのIssueを探しながら行うと良いかもしれません。

作るものとしては、TypeScriptのソースを読み込み、unassertifyした上で別のファイル(もしくは標準出力)にTypeScriptのファイルを書き出すような実装にしようと思います。TypeScript Compiler APIでは本来の機能ではTypeScriptからJavaScriptにtranspileしますが、今回transpileまでは行いません。

unassertifyの実装

早速コードの方を見ていこうと思います。今回のコードは下記に上がっています。src/trans.tsを参照してください。

github.com

他にもやり方はあるのですが、今回は最もシンプルに、ts.createSourceFile,ts.transformを使うやり方を見ていきたいと思います。

import * as ts from "typescript";


// fileを文字列として読み込む
const src = fs.readFileSync(path.resolve(/*path to file*/), {
  encoding: "utf-8"
});

const sourceFile = ts.createSourceFile("", src, ts.ScriptTarget.ES2015);

const result = ts.transform(sourceFile, 
    [removeImport/*ここにtransformerを入れる*/]
);
result.dispose();

const printer = ts.createPrinter();
console.log(printer.printFile(result.transformed[0] as ts.SourceFile));

fs.readFileSyncでファイルを文字列として読み込み、その文字列をts.createSourceFileに渡しています。ts.createSourceFileの第1引数はファイル名の文字列が入るのですが、今回は空で大丈夫です。 createSourceFileの返り値は、Nodeなどの情報が格納されたASTとなっています。 ts.createPrinter()でこのAST(SourceFile)をprint再びTypeScriptのソースに復元することができます。

ts.transformの第1引数にSourceFile, 第2引数にASTを変換するTransformerFactoryを渡すことができます。

このコードがTransoformerFactoryの関数のサンプル実装です。

TransformerFactorの型定義は下記のようになっています。

type TransformerFactory<T extends Node> = (context: TransformationContext) => Transformer<T>;
 
type Transformer<T extends Node> = (node: T) => T;

本質的には、Nodeを受け取り、Nodeを返す関数を生成する関数です。

const removeImport = <T extends ts.Node>(context: ts.TransformationContext) => (
  rootNode: T
) => {
  function visit(node: ts.Node): ts.Node {
    node = ts.visitEachChild(node, visit, context);
    if (!ts.isImportDeclaration(node)) {
      return node;
    }
    const importDecl: ts.ImportDeclaration = node;
    if ((importDecl.moduleSpecifier as any).text === "assert") {
      return null;
    }

    return node;
  }
  return ts.visitNode(rootNode, visit);
};

この関数はassertimportしている行を削除するtransformerです。 Nodeを削除したい場合はnullreturnすると削除することができます。ただし、生成したASTがinvalidな状態だと、思わぬところで実行時例外が発生する可能性があるので、注意が必要です。

このtransformerは、NodeがImportDeclarationかどうかをチェックしています。

    if (!ts.isImportDeclaration(node)) {
      return node;
    }

この部分がちょうどそれに当たります。TypeScriptは、すべてのNodeにkindの数値が割り当てられており、どの種類のNodeなのかを判別することができます。このあたりで実装されています。 kindの数値を覚えていればどのNodeなのかの判別は効きますが、流石につらすぎるので、src/compiler/utilities.tsで定義されているような関数を用いてNodeの種類を判別することをおすすめします。

また、どの行がなんのASTに変換されるのかを意識するのも非常に難しいと思います。ts-ast-viewerなどのツールを使いながら、目的の文や式がなんのASTに当たるのかをチェックしながら実装するのをおすすめします。

この状態でテスト用に下記のファイルを用意します。

import * as assert from "assert";

console.log("hello");

ts-nodeや、transpileなどを用いて、このファイルpathを引数に与えつつ、実装したコマンドを実行してください。

$ node lib/index.js --path {path to example file}
console.log("hello"); # importが消えて入れば成功

ちなみに、削除する場合はnullを返せばよいですが、Nodeを書き換えたい・置き換えたい場合は、引数で与えられるNodeの値を書き換えることにより可能になります。 例えば、assertpower-assertのimportに書き換えたい場合は下記のように実現できます(かなり強引ですが...)。

const swapImport = <T extends ts.Node>(context: ts.TransformationContext) => (
  rootNode: T
) => {
  function visit(node: ts.Node): ts.Node {
    node = ts.visitEachChild(node, visit, context);
    if (!ts.isImportDeclaration(node)) {
      return node;
    }
    const importDecl: ts.ImportDeclaration = node;
    if ((importDecl.moduleSpecifier as any).text === "assert") {
      (importDecl.moduleSpecifier as any).text = "power-assert"; // ここでmoduleSpecifierのStringLiteralを書き換える。
      return node;
    }

    return node;
  }
  return ts.visitNode(rootNode, visit);
};

引き続き、assertcall部分を削除します。transfomerは下記のようになります。

const removeAssertExpression = <T extends ts.Node>(
  context: ts.TransformationContext
) => (rootNode: T) => {
  function visit(node: ts.Node): ts.Node {
    node = ts.visitEachChild(node, visit, context);
    if (!ts.isExpressionStatement(node)) {
      return node;
    }
    if (!ts.isCallExpression(node.expression)) {
      return node;
    }
    const call: ts.CallExpression = node.expression;
    if ((call.expression as any).escapedTxt === "assert") {
      return null;
    }
    if (!ts.isPropertyAccessExpression(call.expression)) {
      return node;
    }
    const propAccess: ts.PropertyAccessExpression = call.expression;

    if ((propAccess.expression as any).escapedText === "assert") {
      return null;
    }

    return node;
  }
  return ts.visitNode(rootNode, visit);
};

assertをcallする場合、assert(value)のように直接callするパターンか、assert.deepEqual(...)のようにexportsされた他の関数を使うパターンがあります。TypeScriptのAST的には、ExpressionStatementexpressionのプロパティが、CallExpressionであり、CallExpressionIdentifierなのか、PropertyAccessExpressionなのかで区別することができます。 このあたりの細かい話は、実際にast-viewerを使ってASTを確認しながら書いていくとわかりやすいと思います。

下記のようなコードを用意して、書いたtransformerをチェックしてみましょう。

import * as assert from "assert";

const obj1 = {
  a: {
    b: 1
  }
};
const obj2 = {
  a: {
    b: 2
  }
};

assert(true);

assert.deepEqual(obj1, obj2);

下記のように出力されるはずです。

const obj1 = {
    a: {
        b: 1
    }
};
const obj2 = {
    a: {
        b: 2
    }
};

非常に簡単ステップでしたが、unassertifyを実現することができました。

まとめ

TypeScriptのCompilerAPIを用いて、ASTの変換・編集を行う方法について書きました。 僕自身もまだ良くわかってない状態ですがTypeScript周りの自動生成ツールを作ろうかと思っていて、その折に触れた技術です。 TypeScriptは、compiler pluginを差し込む口がBabelなどと比べるとかなり制限されており、そのあたりでも差分を感じています。

以上です、ありがとうございました!