Use the TS-compiler to traverse the AST processing of your code

preface

This article is to give a preliminary guide to the students who first understand the ts compiler. After a series of examples from the shallow to the deep, they can master the basic use of THE TS compiler, including AST traversal, transform function writing, expression node creation, etc. At the same time, the principle of TS-Loader and TS-node is briefly analyzed. This article focuses on the compilation of the Transform function.

1. Understand ts Compiler createProgram and transpileModule

These two functions are arguably the two main apis of ts Compiler, both of which convert Typescript code to JavaScript.

transpileModule

The difference is that transpileModule only processes text and does no type checking. This is how it works:

import ts from 'typescript'

// Generate the corresponding JS code directly. Even if the TS type has errors, it can still generate code.
const js = ts.transpileModule('/** Your typescript code **/').outputText
Copy the code

By now, experienced front end students have already guessed the underlying principles of transpileOnly, and you’ve probably thought about the global @ts-ignore before

A minimum implementation of ts-Loader (not guaranteed to work) can be given here:

import ts from 'typescript'

// Simple ts-loader
function loader(this: webpack.LoaderContext<LoaderOptions>, contents: string) {
  const callback = this.async()
  // transpileOnly does type erasure only
  const result = ts.transpileModule(contents, {
    compilerOptions: {},})// Return the generated JS
  callback(null, result.outputText, null)}Copy the code

createProgram

So how does createProgram work? CreateProgram scans typescript files for import and other statements, and iterates through each TS file to compile. It also does type checks, throws type exceptions and terminates the compilation process. It’s a lot more complicated to use than transpileModule:

import ts from 'typescript/lib/typescript'

const compilerOptions = {}
// Create a file read/write service (dependent on nodejs)
const compilerHost = ts.createCompilerHost(compilerOptions)

// Create the compiler
const program = ts.createProgram(['./entry.ts'], compilerOptions, compilerHost)

// Output the file using the file service
program.emit()
Copy the code

This is essentially what the TSC orders the lower level to do. You can see that the TS compiler is separated into a compilerHost, which is used to separate write, read, and so on. Of course, you can hijack the host to intercept the output of the file:

const compilerHost = ts.createCompilerHost(compilerOptions)

// Intercept the file
const originalGetSourceFile = compilerHost.getSourceFile
compilerHost.getSourceFile = fileName= > {
  // You can do something like Webpack Alias here
  console.log(fileName)
  return originalGetSourceFile.call(compilerHost, fileName)
}

// Intercept the written file
compilerHost.writeFile = (fileName, data) = > {
  // data is the compiled js code
  console.log(fileName, data)
}
Copy the code

CompilerHost separates the compiler from the environment. The compiler itself is just a compiler implementation, and operations such as writing files can be used as a host interface.

Use the typescript compiler on the browser

With this foundation of environment separation, it is possible to run TS Compiler! On the browser side. Because TS Compiler is also js and requires only a browser-side compilerHost, typescript officially offers a virtual file package @typescript/ VFS that provides browser-side fs compatibility:

import ts from 'typescript'
import tsvfs from '@typescript/vfs' // Virtual file service
import lzstring from 'lz-string' // A compression algorithm

// Create a context from the CDN that contains the ts lib type library, pulled from the CDN
const fsMap = await tsvfs.createDefaultMapFromCDN(
  compilerOptions,
  ts.version,
  true,
  ts,
  lzstring
)
// You can set a virtual file, file name index.ts, file content second parameter
fsMap.set('index.ts'.'/** typescript code **/')

const system = tsvfs.createSystem(fsMap)
// Host is the part of the TS compiler that isolates file operations
// Here we can create a virtual file service that does not depend on Nodejs and is available in the browser
const host = tsvfs.createVirtualCompilerHost(system, compilerOptions, ts)

// Create the compiler
const program = ts.createProgram({
  rootNames: [...fsMap.keys()],
  options: compilerOptions,
  host: host.compilerHost,
})
Copy the code

Ts -node principle

If you’ve used typescript a lot, you’ve probably used TS-Nodes to execute TS files.

With the above foundation of TS compiled code, ts-Node should do the following:

The TS file is read, converted to JS using transpileModule functions, and eval is performedCopy the code

1. But the eval function is not safe. It can modify anything in the global context. The VM module provides a safer runInThisContext function in nodeJS to control the scope of access.

2. TranspileModule does not process import statements to load modules. Ts-node borrows the module loading mechanism inherent in nodejs and overrides the require function:

import ts from 'typescript'

function registerExtension() {
  // the require function compiles and executes the ts file
  // exports are returned by js execution
  const old = require.extensions['.ts']

  // Define the ts loading function
  require.extensions['.ts'] = function (m: any, filename) {
    // Module's own compilation method
    // We can execute js to get the exports variable, commonJS specification
    const _compile = m._compile

    // the module.prototype. _compile method is used to compile and load js files
    // But the documentation does not indicate that, only look at nodejs source code to know
    // https://github.com/nodejs/node/blob/da0ede1ad55a502a25b4139f58aab3fb1ee3bf3f/lib/internal/modules/cjs/loader.js#L1065
    // https://github.com/nodejs/node/blob/da0ede1ad55a502a25b4139f58aab3fb1ee3bf3f/lib/internal/modules/cjs/loader.js#L1017
    // The underlying principle is runInThisContext
    m._compile = function (code: string, fileName: string) {
      // Compile the TS file using the TS compiler
      const result = ts.transpileModule(code, {
        compilerOptions: {},})// Use the default js compiler function to get the return value
      return _compile.call(this, result, fileName)
    }

    return old(m, filename)
  }
}
Copy the code

2. Know TS Compiler AST Visitor and Transformer

A Visitor is used to traverse the AST tree, and transformer is used to transform the AST tree. The TS compiler API distinguishes forEachChild and visitEachChild, similar to the difference between forEach and map.

Use the visitor to traverse the AST

The AST Node type in TS Compiler is TS.node. To iterate over an AST of ts code, first create a sourceFile:

// Create the ast root
const rootNode = ts.createSourceFile(
  `input.ts`.'/** typescript code **/,
  ts.ScriptTarget.ES2015,
  /*setParentNodes */ true
)
Copy the code

Ts provides a forEachChild function to traverse the children of an AST node:

// Traverse the child nodes of rootNode
ts.forEachChild(rootNode, node= > {
  // The AST node has a kind attribute indicating the type
  console.log(node.kind)

  // The type guard of ts is usually used to distinguish node types:
  // For example, determine whether the current node is an import declaration
  // node: ts.Node
  if (ts.isImportDeclaration(node)) {
    // node: ts.ImportDeclaration
    // There is a getText method on the node that prints the text of the node for debugging purposes
    console.log(node.getText())
  }
})
Copy the code

ForEachChild only traverses down one level, but we need to traverse the entire ast tree, so we need to recursively traverse it:

const traverse = (node: ts.Node) = > {
  console.log(node.getText())
  // Iterate through each node recursively
  ts.forEachChild(node, node= > traverse(node))
}
Copy the code

With the base of that visitor, you can implement the import statement for the analysis code:

const traverse = (node: ts.Node) = > {
  if (ts.isImportDeclaration(node)) {
    // The imported package name or path is the moduleSpecifier module identifier
    const library = node.moduleSpecifier.getText()

    // Import by default
    constdefaultImport = node.importClause? .name?.getText()// Deconstruct the import
    constbindings = node.importClause? .namedBindingsas ts.NamedImports
    // names is the name of the deconstructed variable
    const names = bindings.elements.map(item= > item.getText())
  }
  // Iterate recursively
  ts.forEachChild(node, node= > traverse(node))
}
Copy the code

If you need to import all the code under the file, you can use the @nodelib/fs.walk library to get all the files.

Use Transformer to transform the AST

The ts compiler provides a transform interface to modify the AST, createProgram, and transpileModule transform interfaces as follows:

// createProgram
const program = ts.createProgram(['input.ts'], options, compilerHost)
// The last parameter is CustomTransformers
program.emit(undefined.undefined.undefined.undefined, transformers)

// transpileModule
// The second parameter has a transformers interface
ts.transpileModule('/** typescript code **/, { transformers })
Copy the code

Transformers is an object that provides three lifecycle hooks: before, at, and after compilation:

ts.transpileModule('/** typescript code **/, {
  transformers: {
    before: [].// Before transformer can obtain the ts type and operate on the TS type ast
    afterDeclarations: [].after: [].// Transformer in after does not have ts type and can only operate js code AST}})Copy the code

With the transformers interface, you can write a “simple” Transformer. First, the visitEachChild function provides ast modification capability:

export function visitNodes(node: ts.Node) {
  // Recursively, the second function argument returns a value of type TS.node, replacing the AST Node
  Array.prototype.map
  return ts.visitEachChild(node, childNode= > visitNodes(childNode))
}

ts.transpileModule('/** typescript code **/, {
  transformers: {
    // The transform function is a high-order function
    // transform :: TransformationContext -> ts.Node -> ts.Node
    before: [context= > node= > visitNodes(node)],
  },
})
Copy the code

In the visitNodes recursive function, which determines the node type and returns a new node replacement, a “minor requirement” can be implemented as follows:

The function adds a jsDoc annotation @noexcept that automatically adds a try and catch during static compilation, like this:

/ * * *@noexcept* /
function main () :number {
  throw new Error()
}

main() // There will be no exception, the error will be logged out
Copy the code

First, we need to determine whether the node is a function declaration node:

// isFunctionDeclaration determines that the current node is a function declaration
if (ts.isFunctionDeclaration(node)) {
  console.log(node.getText()) // Print the complete function definition, including the jsDoc comment
}
Copy the code

Then get the contents of jsDoc and determine if there is @noexcept:

if (ts.isFunctionDeclaration(node)) {
  // The getJSDocTags function gets tags from jsDoc
  constenableNoexcept = !! ts.getJSDocTags(node).find(tag= > {
    // Determine if there is a @noexcept annotation in the jsDoc annotation
    return tag.tagName.escapedText === 'noexcept'})}Copy the code

The judgment function uses the @noexcept annotation, and now you can make changes to the AST node (replace it with a new one).

To wrap the function body around a trycatch, first get the function body:

if (ts.isFunctionDeclaration(node)) {
  // Properties of function nodes
  node.decorators // Function decorator
  node.modifiers // Don't know what it is
  node.asteriskToken // Don't know what it is
  node.name / / the function name
  node.typeParameters // Function type parameters
  node.parameters // Function parameters
  node.type // Function type
  node.body // This is what we need

  // Example: you can create an identical function declaring node Clone:
  return ts.factory.createFunctionDeclaration(
    node.decorators,
    node.modifiers,
    node.asteriskToken,
    node.name,
    node.typeParameters,
    node.parameters,
    node.type,
    node.body // Add a try and catch statement)}Copy the code

Next, create a trycatch statement using node.body as the contents of the try block

// Create a try catch statement
const tryStatement = ts.factory.createTryStatement(
  node.body, // The contents of the try block. Node. body is the body of the function that needs to be wrapped. Let's create the catch part
  // Create a catch
  ts.factory.createCatchClause(
    'error'.// Catch parameter name
    // Catch the content of the block
    ts.factory.createBlock([])
  ),
  undefined // Finally does not create, only need to handle catch
)
/** * The above code does the following: * try {* body *} catch (error) {* error handling *} */
Copy the code

However, catch statements usually need to handle exceptions, otherwise errors can be swallowed and hidden. Here we can add a console.log call statement to the catch:

// Create an expression statement
const consoleErrorStatement = ts.factory.createExpressionStatement(
  // Create a function call expression
  ts.factory.createCallExpression(
    // Create the call statement, console.log(error)
    ts.factory.createIdentifier('console.log'), // Create an identifier (function call)[].// Type parameter, none
    [ts.factory.createIdentifier('error')] // Pass in parameters))Copy the code

With the try catch statement and the consoleError statement, the following can be combined, complete code:

// This is a simple transform function
export const transformNoExcept = node= > {
  if (ts.isFunctionDeclaration(node)) {
    constenable = !! ts.getJSDocTags(node).find(tag= > {
      return tag.tagName.escapedText === 'noexcept'
    })
    if (enable) {
      return ts.factory.createFunctionDeclaration(
        node.decorators,
        node.modifiers,
        node.asteriskToken,
        node.name,
        node.typeParameters,
        node.parameters,
        node.type,
        ts.factory.createBlock([
          ts.factory.createTryStatement(
            node.body,
            ts.factory.createCatchClause(
              'error'. ts.factory.createBlock([ ts.factory.createExpressionStatement( ts.factory.createCallExpression( ts.factory.createIdentifier('console.log'),
                    [],
                    [ts.factory.createIdentifier('error'))))))undefined),]))}}return node
}
Copy the code

For transpile use:

export function visitNodes(node: ts.Node) {
  // Handle the @noexcept comment and add tryccatch
  const newNode = transformNoExcept(node)
  if(node ! == newNode) {return newNode
  }
  return ts.visitEachChild(node, childNode= > visitNodes(childNode))
}

ts.transpileModule('/** typescript code **/, {
  transformers: {
    before: [context= > node= > visitNodes(node)],
  },
})
Copy the code

Ts-loader uses the transform function

You’ve seen how to write a transform function, but how does it fit into a real project? The Ts Compiler API exposes the Transform interface, while TSC and TSConfig do not. Ts-loader exposes the transform interface:

// webpack.config.js
module.exports = {
  module: {
    rules: [{test: /\.ts$/,
        loader: 'ts-loader'.options: {
          // The return value of getCustomTransformers is transformers
          getCustomTransformers: () = > ({
            // You can set your own transform function
            before: [context= > node= > node],
          }),
        },
      },
    ],
  },
}
Copy the code

conclusion

1. Using TS-Compiler allows you to analyze code more accurately, such as extracting imports. If you write your own regular expressions to match the code, it’s easy to miss cases, and the syntax can vary with language standards.

2. Transform plug-in functions can be written for TS-Compiler to achieve custom syntax sugar.

---

The above code is open source on Github:

ts-compiler