Custom Babel Plugin to Sneak Code in

Babel is source-to-source JavaScript compiler (or transpiler), simply means it takes one piece of source code and transform to another piece that can run in targeting platforms (specific browser versions or NodeJS). Compared with those traditional compiler, it still produces source code as output.

The steps of transpiling

We can see the whole process is to eventually return another piece of source of code. If there is no plugin applied, there is no transformation occurred. The output is the same as input source code.

Plugin as the unit of transformation

As shown in previous transpiling steps, each plugin performs a specific transformation. Preset is a set of plugins pre-built for us, so we don't need to include one by one.
If we work on modern JS projects, chances are we've already used them either by configuring them manually or through 3rd party tooling packages or frameworks (CRA, Next.js). These are major ones provided by Babel.

The ingredients of a plugin

All plugins work at 'transform' phase, which manipulate AST nodes they are interested in while Babel is traversing. Visitor is the way that plugins know which node Babel is visiting, register a function to do the work.
Visitors let a plugin know where the transformation should happen, another concept is Path to actually do the transformation: add, replace or remove the nodes

Each plugin is essentially a visitor, use Path to access info associated with the node, and mutate them. Every plugin will have the following format:

const MyPlugin =  {
  return {
    visitor: {
       NodeType(path, state)  {
          // do traversal, transform etc. for this node
       }
    }
  }
}

Use Cases

A common one is to inject HOCs to wrap React components for some additional logic, like sending some analytic events, or instrument some profiling logic, it still returns same decorated component, from outside we don't need change comp hierarchy. This is suitable to use Babel to inject these wrapping, some benefits:

No need to manually wrap every single component, and
We don't accidentally commit these code, less intrusive to the source code.
Integrate with pipelines. Use babel.config.js to control which env we need to inject and could automate it during building process.

Assume the component defined like this

export const MyComp = (props) => {
  // does not matter what it does
  return (
    <div>
       // does not matter what it returns
    </div>
  )
}

And the goal is to output source code like below:

// this import needs to be inserted dynamically to the original import list
+ import withSomeMagic from './hoc'
+ export const MyComp = withSomeMagic((props) => {
   return (
    <div>
       // does not matter what it returns
    </div>
  )
})

So some magic code will be injected automatically without changing the way MyComp is used.

Before jumping into the code, we could use Ast Explorer to help visualise the tree structure (AST), so we have better idea what the code is doing. As shown in the screenshot, from high level, the whole module consists of ImportDeclaration and ExportNamedDeclaration

The whole export is represented as ExportNamedDeclaration node, the Function component itself is a VariableDeclarator node.
VariableDeclarator contains id for exported component name, which has type Identifier
It also contain init for arrow function definition, which has type ArrowFunctionExpression. This is what we use to refer to the component function body, and get it wrapped.

The Fun part

With the AST structure, now we can try to figure out how to traverse and manipulate them with Babel to output our expected code.

const t = require('@babel/types')
const template = require('@babel/template')

const MyPlugin = () => {
  return {
    visitor: {
       // the first visitor is to manipulate the 'import' list
       Program(path, state) {
         // insert import HOC
         const idf = t.identifier('withSomeMagic')
         const importSpec = t.importSpecifier(idf, idf)
         // stringLiteral is the path where hoc is imported
         const importDeclaration = t.importDeclaration([importSpec], t.stringLiteral('./hoc'))

         path.unshiftContainer("body", importDeclaration);
       },

       ExportNamedDeclaration(path, state) {
         // file name contains full path for current file being transpiled
          const filename = state.file.opts.filename
          try {
            console.log('## ExportNamedDeclare.visit', filename);
            const def = path.get('declaration.declarations.0')

            const compName = def.node.id.name
            const origArrowFn = def.node.init
           // template to easily build new AST with our HoC wrapping the component function
            const wrappedFn = buildHOCWrapper({
              ORIGINAL_FN_CALL: origArrowFn,
            })
             // real magic happens here, replace the original `init` node
            path.get('declaration.declarations.0.init').replaceWith(wrappedFn)
          } catch (e) {
            console.log(`## ExportNamedDeclare.visit fail ${filename}`, e.message);
          }
       },
       }
    }
  }
}
const buildHOCWrapper = template.default(`{
  withSomeMagic(ORIGINAL_FN_CALL,)
}`);

module.exports = MyPlugin

A few notes regarding the code:

▪︎ Program visitor
This is mainly to insert import withSomeMagic from './hoc' to original import list.

▪︎ ExportNamedDeclaration visitor This is where we wrap the original functional component with the high order function. Note Babel will try to apply this plugin in every module which has export const XYZ =, we could try to limit it on React function component only.
We could use path.traverse inside ExportNamedDeclaration to first check its return by providing another visitor to it. We do early return if the return does not meet.

        const checkReturn = {
          ReturnStatement(path) {
            // this points to `traverseStatus` object passed in
            this.isReactComp = path.node.argument && path.node.argument.type === 'JSXElement'
          }
        }
        const traverseStatus = {
          isReactComp: false
        }
        path.traverse(checkReturn, traverseStatus)
        if (!traverseStatus.isReactComp) return

▪︎ Config babel.config.js In order to use the plugin, we need to config it in babel.config.js. We could also provide some options in plugin config, e.g could give some paths where plugin only applied on modules contained in any of the paths.

    module.exports = {
      ['./my-babel-plugin.js',
        {
          filePaths: [
            'to/comp/path1',
            'to/comp/path2',
          ]
        }
      ],
    }

To get those options, use state.opts.filePaths. If we install babel cli, run npx babel {input.file} -o {output.file} to see if plugin takes effect.

▪︎ Use babel-template This allows to easily build AST nodes from string (representing some chunk of source code)

With everything coded up, when it starts bundling the source code, all config plugins will kick in, and we could see some code magically get injected in the bundled code.🤠

6thfdwp / learning-thoughts