JavaScript AST Manipulation for Fun and Profit
For many JavaScript developers, the code they write is the final product. But what if you could write code that understands and even rewrites other code? This isn't science fiction; it's the power of Abstract Syntax Trees (ASTs). By the end of this article, you'll understand what ASTs are, how to manipulate them, and how to build your own codemods to automate refactoring, and boost your productivity.
What is an Abstract Syntax Tree (AST)?
At its core, an AST is a tree representation of your code. Think of it like a diagram of a sentence, but for code. Each part of your code, from variables and functions to expressions and statements, is represented as a node in the tree. This tree structure provides a powerful way to analyze and manipulate code programmatically.
Let's take a simple example:
const sum = 1 + 2;
This line of code can be represented as the following AST:
{
"type": "Program",
"body": [
{
"type": "VariableDeclaration",
"declarations": [
{
"type": "VariableDeclarator",
"id": {
"type": "Identifier",
"name": "sum"
},
"init": {
"type": "BinaryExpression",
"left": {
"type": "NumericLiteral",
"value": 1
},
"operator": "+",
"right": {
"type": "NumericLiteral",
"value": 2
}
}
}
],
"kind": "const"
}
]
}
As you can see, the AST provides a detailed, structured representation of the code. This structure is what allows us to traverse and manipulate the code with precision.
To explore ASTs further, I highly recommend using AST Explorer. It's an online tool that lets you write code and see its corresponding AST in real-time. It's an invaluable resource for learning and experimenting with ASTs.
Prerequisites and Environment Setup
Before we dive into the world of AST manipulation, you'll need a few things:
- Node.js and npm/yarn: If you don't have them installed, you can download them from the official Node.js website.
- A code editor: I recommend using Visual Studio Code, but any code editor will do.
Once you have the prerequisites, create a new directory for our project and initialize a new Node.js project:
mkdir js-ast-manipulation
cd js-ast-manipulation
npm init -y
Next, we'll need to install the necessary Babel packages. Babel is a popular JavaScript compiler that provides a suite of tools for working with ASTs.
npm install --save-dev @babel/core @babel/parser @babel/traverse @babel/generator
Here's a breakdown of what each package does:
@babel/core: The main Babel package.@babel/parser: Parses JavaScript code and generates an AST.@babel/traverse: Provides a simple and powerful way to traverse and manipulate the AST.@babel/generator: Generates JavaScript code from an AST.
Parsing: From Code to AST
The first step in any AST manipulation task is to parse the code and generate an AST. The @babel/parser package makes this easy.
Let's create a file named parser.js and add the following code:
const parser = require("@babel/parser");
const code = "const sum = 1 + 2;";
const ast = parser.parse(code);
console.log(JSON.stringify(ast, null, 2));
Now, run the script from your terminal:
node parser.js
You should see the same AST that we saw earlier. The @babel/parser package has taken our string of JavaScript code and transformed it into a detailed, structured tree that we can now traverse and manipulate.
Traversal: The Heart of AST Manipulation
Now that we have an AST, we can start to traverse it. Traversal is the process of visiting each node in the tree. This is where the @babel/traverse package comes in.
@babel/traverse uses the Visitor pattern to traverse the AST. The Visitor pattern is a design pattern that allows you to add new behavior to an existing object structure without modifying the structure itself. In our case, it allows us to define a set of "visitor" functions that will be called for each node in the AST.
Let's create a file named traverse.js and add the following code:
const parser = require("@babel/parser");
const traverse = require("@babel/traverse").default;
const code = "const sum = 1 + 2;";
const ast = parser.parse(code);
traverse(ast, {
enter(path) {
console.log(`Entering: ${path.node.type}`);
},
exit(path) {
console.log(`Exiting: ${path.node.type}`);
}
});
When you run this script, you'll see a list of all the nodes in the AST, in the order that they are visited. The enter function is called when the traversal enters a node, and the exit function is called when it exits a node.
The path object that is passed to our visitor functions is a wrapper around each node that provides a wealth of information about the node, including its parent, its siblings, and its scope.
Manipulation: Modifying the AST
Now for the exciting part: manipulating the AST. With @babel/traverse, we can not only visit each node in the tree, but we can also modify them.
Let's say we want to rename the sum variable to total. We can do this by modifying the name property of the Identifier node.
Create a file named manipulate.js and add the following code:
const parser = require("@babel/parser");
const traverse = require("@babel/traverse").default;
const generate = require("@babel/generator").default;
const code = "const sum = 1 + 2;";
const ast = parser.parse(code);
traverse(ast, {
Identifier(path) {
if (path.node.name === "sum") {
path.node.name = "total";
}
}
});
const output = generate(ast, code);
console.log(output.code);
When you run this script, you'll see the following output:
const total = 1 + 2;
We've successfully renamed the sum variable to total!
Generation: From AST back to Code
The final step in our AST manipulation journey is to generate human-readable code from our modified AST. The @babel/generator package makes this a breeze.
As you saw in the previous example, the generate function takes our modified AST and returns an object with a code property that contains the generated code.
The generate function also takes an optional options object that allows you to customize the output. For example, you can use the retainLines option to preserve the original line numbers, or the compact option to generate minified code.
Building a Practical Codemod: A Step-by-Step Guide
Now that we've covered the basics of AST manipulation, let's build a practical codemod that automatically replaces all instances of var with let or const.
1. Set up the project
Create a new file named codemod.js and add the following code:
const parser = require("@babel/parser");
const traverse = require("@babel/traverse").default;
const generate = require("@babel/generator").default;
const code = `
var a = 1;
var b = 2;
b = 3;
`;
const ast = parser.parse(code);
// ...
2. Traverse the AST and replace var with let or const
Next, we'll traverse the AST and replace all VariableDeclaration nodes that have a kind of var with let or const. We'll use the scope property of the path object to determine if a variable is ever reassigned. If it is, we'll use let. Otherwise, we'll use const.
traverse(ast, {
VariableDeclaration(path) {
if (path.node.kind === "var") {
const isReassigned = path.node.declarations.some(declaration => {
const binding = path.scope.getBinding(declaration.id.name);
return binding && binding.constantViolations.length > 0;
});
path.node.kind = isReassigned ? "let" : "const";
}
}
});
3. Generate the transformed code
Finally, we'll generate the transformed code from our modified AST:
const output = generate(ast, code);
console.log(output.code);
When you run this script, you'll see the following output:
const a = 1;
let b = 2;
b = 3;
Our codemod has successfully replaced var with let and const!
Conclusion
In this article, we've only scratched the surface of what's possible with AST manipulation. From automating complex refactoring to creating your own custom language features, the possibilities are endless.
You now have the knowledge and tools to start your own AST manipulation journey. I encourage you to experiment with the concepts we've covered today and to explore the wealth of resources available online.
