Search code examples
typescripttypescript-compiler-api

How can I inject additional statements into a function using ts.transform


I'm using the Typescript compiler API (ts.transform, ts.updateFunctionDeclaration) to inject additional statements at the beginning of functions in existing source files. This works great, except that when I print the transformed code (using ts.Printer) the first comment in the original source is emitted above my injected statements. How can I prevent that happening?

I've tried various ways of constructing the new nodes returned by the transformer, but nothing changes the output. Here's my complete code, I'm executing it using Node.

    import * as ts from 'typescript';

    const sourceLines = [
        "// The list of properties exposed by this script to the user",
        "var properties = {'speed': 20};",
        "",
        "function tick(deltaTime)",
        "{",
        "  var currentDeltaTime = deltaTime;",
        "}"
    ];

    var sourceCode = sourceLines.join('\n');

    const sourceFile = ts.createSourceFile("originalSource.js", sourceCode, ts.ScriptTarget.Latest, true);

    const additionalSource: ts.SourceFile = ts.createSourceFile(
        'ownerCheck.js', "var variableToAdd = 5;", ts.ScriptTarget.ES5, false, ts.ScriptKind.JS
    );

    const transformer = <T extends ts.Node>(context: ts.TransformationContext) =>
        (rootNode: T) => {
            function visit(node: ts.Node): ts.Node {

                if (ts.isFunctionDeclaration(node)) {
                    var functionDeclaration = <ts.FunctionDeclaration>node;
                    var functionName: ts.Identifier = functionDeclaration.name;

                    var updatedFunction = ts.updateFunctionDeclaration(
                        functionDeclaration,
                        /*decorators*/     undefined,
                        /*modifiers*/      undefined,
                        /*asteriskToken*/  functionDeclaration.asteriskToken,
                        /*functionName*/   functionName,
                        /*typeParameters*/ undefined,
                        /*parameters*/     functionDeclaration.parameters,
                        /*returnType*/     undefined,
                        /*body*/           ts.createBlock(ts.createNodeArray([].concat(additionalSource.statements, functionDeclaration.body.statements), false))
                    );

                    return updatedFunction;

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

    const result = ts.transform(
        sourceFile, [transformer]
    );

    const transformedNodes = result.transformed[0];

    const printer: ts.Printer = ts.createPrinter({
        newLine: ts.NewLineKind.LineFeed,
        removeComments: false
    });

    console.log(printer.printNode(ts.EmitHint.SourceFile, transformedNodes, sourceFile));

Output:

    // The list of properties exposed by this script to the user
    var properties = { 'speed': 20 };
    function tick(deltaTime) {
        // The list of properties exposed by this script to the user
        var variableToAdd = 5;
        var currentDeltaTime = deltaTime;
    }

I would expect the output to include the additional statements without the repeated comment prefixing them.

Update: Adding this function:

stripRanges(additionalSource);

function stripRanges(node: ts.Node) {
    node.pos = -1;
    node.end = -1;

    ts.forEachChild(node, stripRanges);
}

to strip the node ranges in the additionalSource nodes does technically work, but removes all linefeeds in the output after those nodes resulting in the following output:

// The list of properties exposed by this script to the user
var properties = { 'speed': 20 };
function tick(deltaTime) { var variableToAdd = 5; var currentDeltaTime = deltaTime; }

Update 2: Changing the ts.createBlock() call to have multiLine set to true fixed the output:

ts.createBlock(
    ts.createNodeArray(
        [].concat(
            additionalSource.statements,
            functionDeclaration.body.statements),
    false), /*has trailing comma*/ 
true) /*multiLine*/ 

Solution

  • This is because when the printer goes to print the additionalSource.statements nodes it is using the sourceFile node's text to get the comments based on the position of the nodes in additionalSource.

    You can see this reproduced by running the following code:

    console.log(ts.createPrinter().printNode(
        ts.EmitHint.Unspecified,
        additionalSource.statements[0],
        sourceFile // it's getting the comments from this sourceFile.text
    ));
    

    Outputs:

    // The list of properties exposed by this script to the user
    var variableToAdd = 5;
    

    Solution

    A workaround is to remove the positions from the nodes in additionalSource before using it in sourceFile:

    stripRanges(additionalSource);
    
    function stripRanges(node: ts.Node) {
        node.pos = -1;
        node.end = -1;
    
        ts.forEachChild(node, stripRanges);
    }
    

    I'm using -1 for the pos and end because that's what the compiler does when a node is made with a factory function. For example:

    const binaryExpr = ts.createBinary(1, "+", 2);
    binaryExpr.pos; // -1
    binaryExpr.end; // -1
    

    Side note: The way the TypeScript compiler deals with comments is not ideal in my opinion... especially for transformation and printing (I prefer how Babel deals with comments, but there are improvements that could be made there too...).