Modifying JavaScript AST with Yeoman

Yeoman helped teams scaffold projects and modules faster and more comfortable than ever. Adding new files to a project became stupidly simple. Files were treated as templates, feeded with data and copied around. Updating files was a bit of a different story though.

Yeoman and babylon — ES2015 all the way. (yeoman.io / babeljs.io)

When Yeoman was released JavaScript files that needed to be changed were treated as that: files. Often people added markers like /* ADD IMPORT */, which were then matched with Regular Expressions and replaced by a string. If you have ever attended Theoretical Computer Science classes you might also remember that parsing a program with RegEx is not reliable. The latest Yeoman documentation advise you to parse the file’s AST and modify it, rather than parsing it with RegEx.

Parsing a code file with RegEx is perilous path, and before doing so, you should read this CS anthropological answers and grasp the flaws of RegEx parsing.

The StackOverflow reply is a classic, and is definitely worth a read through. The documentation also states the following:

Updating a pre-existing file is not always a simple task. The most reliable way to do so is to parse the file AST (abstract syntax tree) and edit it. The main issue with this solution is that editing an AST can be verbose and a bit hard to grasp.

The Challenge

Although the docs make it obvious what we should do to change files, they don’t explain how the AST can be parsed and modified.

Recently I’ve been working on a small generator for a React + Redux project. One of the sub-generators should add an action, a reducer and import the created reducer in the root reducer, which looks like this:

import { combineReducers } from ‘redux’;
import examplesReducer from ‘./reducers/examples’;
export default combineReducer({
  examples: examplesReducer,
});

For this article I’ll focus on how to automatically import the newly created reducer, but adding it to the argument provided to combineReducer follows the same approach.

From Text to AST

To work with a JavaScript AST you have a couple of options, three of which are quite popular: Babel, TypeScript and Esprima. Comparing these is a topic of its own and would go beyond the scope of this article. Because of my familiarity with Babel and some neat libraries around it, especially regarding AST traversal, I chose to go with it.

Parsing the AST from a file is actually quite easy:

const babylon = require('babylon');
const source = this.fs.read(
  this.destionationPath('src/reducers.js')
);
const ast = babylon.parse(source, {
  sourceType: 'module',
});

A (very!) simplified version of the AST of the reducer file above would look like this:

[
  {
    "type":"ImportDeclaration",
    "specifiers":[
      {
        "type":"ImportSpecifier",
        "imported":{
          "name":"combineReducers"
        },
        "source":{
          "type":"StringLiteral",
          "value":"redux"
        }
      }
    ]
  },
  {
    "type":"VariableDeclaration",
    "declarations":[
      {
        "type":"VariableDeclarator",
        "id":{
          "type":"Identifier",
          "name":"rootReducer"
        },
        "init":{
          "type":"CallExpression",
          "callee":{
            "type":"Identifier",
            "name":"combineReducers",
            "arguments":[...]
          }
        }
      }
    ],
    "kind":"const"
  }
]

These two nodes roughly represent these two lines:

import { combineReducers } from 'redux';
const rootReducer = combineReducers(/* ... */);

Every single piece of code is represented in the AST. To modify our code, we need to traverse along the tree and find the correct node we want to modify.

Adding an Import Declaration

Since we already import from the redux package, but don’t know how many reducers are already being imported, we want to add the new declaration after the last one.

The babel-traverse package helps us walk the AST, with babel-types providing helpers to identify certain tokens. A good way to travel trees is the visitor design pattern, which babel-traverse leverages.

const traverse = require('babel-traverse').default;
let lastImport = null;
traverse(ast, {
  ImportDeclaration(path) {
    lastImport = path;
  }
});

Adding the new import declaration after the last one means it should be added when the next token is not another import. Unfortunately the pattern doesn’t offer us a way to look ahead in the tree and see what the next node will be. Therefore we need to remember the last import declaration.

If we let this run and inspect lastImport you’ll see that it actually is the last import declaration. If we visit a node we can check whether lastImport was defined and the currently visited node is not another import declaration.

traverse(ast, {
  enter(path) {
    if (lastImport && !isImportDeclaration(path)) {
      // Add the new import declaration
    }
  }
});

In a real AST the import declaration is not just one node, therefore isImportDeclaration is a tad more complicated.

const t = require('babel-types');
const isImportDeclaration = path => 
  t.isImportDeclaration(path.node) ||
  t.isImportSpecifier(path.parent) ||
  t.isImportDeclaration(path.parent) ||
  t.isImportSpecifier(path.parent) ||
  t.isImportDefaultSpecifier(path.parent);

At this point we know 1) that the current node is not another import declaration and 2) the last import declaration. With that information we can insert a new declaration.

const declaration = t.importDeclaration(
  [t.importDefaultSpecifier(t.identifier('tasksReducer'))],
  t.stringLiteral('./reducers/tasks'),
);

t.importDeclaration assigns a new node to declaration, which will be resolved in code generation to import tasksReducer from './reducers/tasks';. Fortunately adding declaration to the AST is fairly easy.

traverse(ast, {
  enter(path) {
    if (lastImport && !isImportDeclaration(path)) {
      lastImport.insertAfter(declaration);
    }
  },
  ImportDeclaration(path) {
    lastImport = path;
  }
});

🎉 Congratulations, the hard work is done! Now you only need to transform the AST to code again, which is incredibly simple with babel-generator.

const generate = require('babel-generator').default;
const { code } = generate(ast, { /* config */ }, source);
this.fs.write(this.destinationPath('src/reducers.js', code);

That’s it! Of course this pattern is adaptable to other AST modifications as well. For example you could automatically map the state key to its reducer by implementing an ObjectProperty visitor. But honestly, by having access to the AST you can do anything you’d like.

I hope this article helped you to properly modify your JavaScript files with Yeoman. It might all be a bit much just for an import declaration but it’s fun to play around with Babel’s AST and offers you some insights into the language.

For the record, here’s the full code used in a generator: