TypeScriptのASTを使ってコードを書き換えよう

TypeScriptのASTは、コードを更新する機能が用意されています。これを使えば、やりたい放題なんでもできます。ただし、まだ実装が不十分で複数回コードを書き換えると、更新前のASTを参照してしまい壊れていくことがあるので、注意してください。

準備

TypeScriptのASTを使ってコードを解析しよう を参照してください。

ASTの書き換え方

TypeScriptのASTを使ってコードを解析しようのAST解析の進め方 と4までは同じです。Visitor パターンを使って解析しながら書き換えていきます。

ASTを書き換えるには、下位のようなメソッドを用意します

function changeText(updateText: string, node: ts.Node):ts.SourceFile {
    // ASTのノードの開始地点
    var start = node.getFullStart();
    // ASTのノードの終了地点
    var end = node.getEnd();
    // 元のソースコード
    var oldText = sourceFile.text;
    // 書き換える領域の前のコード
    var pre = oldText.substring(0, start);
    // 書き換える領域の後のコード
    var post = oldText.substring(end);
    // 書き換え後のコード全体
    var newText = pre + updateText + post;
    // 書き換える領域
    var textChangeRange: ts.TextChangeRange = {
        span: {
            start: start,
            length: (end - start)
        },

        newLength: (updateText.length)
    }
    // コードを書き換える
    return sourceFile.update(newText, textChangeRange);

}

コードを書き換えるメソッドがTypeScriptのASTの機能(SourceFile#update)が用意されているのですが、これが使いにくい。書き換え後のソースコード全体と書き換える領域を引数で渡す必要があり、２つの引数間の整合性を取らなければならない。上記のchangeTextメソッドで整合性が取れるので、簡単に書き換えることができるようになります。

メソッド名を文字列として書き出してみる

下記のコードで、メソッドの定義の後にメソッド名をテキストとしてstatic変数で書き出しています。

function methodDeclaration(node: ts.MethodDeclaration) {
    // 元のテキストを取得する
    var text = node.getFullText(sourceFile);
    // 書き換え後のコードを組み立てる
    var updateText = text + 'static ' + (<ts.Identifier>node.name).text + '$name = \'' + (<ts.Identifier>node.name).text + '\';'
    // コードを書き換える
    sourceFile = changeText(updateText, node);
    return false;
}

methodDeclarationメソッドをVisitorパターンで呼び出す

Visitorパターンを使って、ASTのノードがメソッド定義だった場合に、コードの書き換えメソッドを呼び出します。

// ソースコードをASTに変換する
let sourceFile = ts.createSourceFile('sample.ts', source, ts.ScriptTarget.ES6, /*setParentNodes */ true);

// Visitorパターンを呼び出す
ts.forEachChild(sourceFile, each);
// 変換後のコードを表示する
console.log(sourceFile.text);

// Visitorパターンでコールバックされるメソッド
function each(node: ts.Node): boolean {
    switch (node.kind) {
        // メソッド定義の場合には、コードを書き換える
        case ts.SyntaxKind.MethodDeclaration:
            return methodDeclaration(<ts.MethodDeclaration>node);

    }

    return next();

    function next(): boolean {
        return ts.forEachChild(node, each);
    }

}

実行例

このコードを動かすと、メソッド宣言の後にメソッドの文字列を保持する関数が追加されます。

変換前

 class Sample{
   constructor(){
   }

   method1(){
   }

   method2(){
   }

   method3(){
   }

 }

変換後

 class Sample{
   constructor(){
   }

   method1(){
   }static method1$name = 'method1';

   method2(){
   }static method2$name = 'method2';

   method3(){
   }static method3$name = 'method3';

 }

まとめ

TypeScriptのASTのコード更新機能をを使うことで、ASTを解析しながらソースコードを書き換えていくことができます。TypeScriptのコード上の情報はすべて残っているので、TypeScriptのコードを自由に書き換えることができます。ただし、上記のコードのようにメソッドの書き換えならうまく動作しましたが、クラスの書き換えた場合には、クラスが２つ以上あるファイルを書き換えると壊れていきました。TypeScript 1.5 Betaの実装では、まだ安定していないメソッドのようなので注意してください。

作成したコード

import * as ts from 'typescript'

let source = `
 class Sample{
   constructor(){
   }

   method1(){
   }

   method2(){
   }

   method3(){
   }

 }

`;

let sourceFile = ts.createSourceFile('sample.ts', source, ts.ScriptTarget.ES6, /*setParentNodes */ true);

ts.forEachChild(sourceFile, each);
console.log(sourceFile.text);

function each(node: ts.Node): boolean {
    switch (node.kind) {
        case ts.SyntaxKind.MethodDeclaration:
            return methodDeclaration(<ts.MethodDeclaration>node);

    }
    return next();

    function next(): boolean {
        return ts.forEachChild(node, each);
    }

}

function methodDeclaration(node: ts.MethodDeclaration) {
    var text = node.getFullText(sourceFile);
    var updateText = text + 'static ' + (<ts.Identifier>node.name).text + '$name = \'' + (<ts.Identifier>node.name).text + '\';'
    sourceFile = changeText(updateText, node);
    return false;
}

function changeText(updateText: string, node: ts.Node):ts.SourceFile {

    var start = node.getFullStart();
    var end = node.getEnd();
    var oldText = sourceFile.text;
    var pre = oldText.substring(0, start);

    var post = oldText.substring(end);
    var newText = pre + updateText + post;

    var textChangeRange: ts.TextChangeRange = {
        span: {
            start: start,
            length: (end - start)
        },

        newLength: (updateText.length)
    }

    return sourceFile.update(newText, textChangeRange);

}