Writing Javascript Codemods and Understanding AST Easily
Tags: javascript, codemod, ast, jscodeshift
Published: 2019-03-05 20:36:12

---
title: Writing Javascript Codemods and Understanding AST Easily
published: true
description: Analysing and modifying abstract syntax tree is not that hard, if you start from simple cases and gradually increase complexity. 
tags: Javascript, codemod, AST, jscodeshift
cover_image: https://thepracticaldev.s3.amazonaws.com/i/yrs4wqg4k9k5gve5w0om.jpg
canonical_url: https://katilius.dev/writing-js-codemods/
---
> Originally posted in my [personal blog](https://katilius.dev/writing-js-codemods/)


One of the big advantages when using statically typed language is ease of refactoring. Different IDE tools can easily rename class or method across hundreds of files with hundreds of usages. And given the nature of Javascript, some of the refactorings are hard, or even impossible.

Despite that, different tools, that modify or inspect Javascript code still emerge. And in some cases, they are even better than ones in statically typed languages ecosystem. **Prettier**, **Eslint**, **React Codemods** to name a few.

They all have one in common - they all analyze or modify parsed Abstract Syntax Tree of the code. Basically, AST let you traverse source code using a tree structure. AST is a general programming languages term and not specific to Javascript. I will not go into the theory about AST here, but I will show a concrete example of how to use it.


## Notable tools and libraries

- [AST Explorer](https://astexplorer.net/) - one of the most useful tools while learning. You paste JS code and see AST representation in different AST specs.
- [jscodeshift](https://github.com/facebook/jscodeshift) - a tool by Facebook, that helps writing code modification scripts.
- [AST Types](https://github.com/benjamn/ast-types/) - type specification on which jscodeshift is based.
- [react-codemod](https://github.com/reactjs/react-codemod) - collection of scripts, written for jscodeshift that convert React code in different ways. There are some good examples to look into.
- [js-codemod](https://github.com/cpojer/js-codemod) - Similar collection of scripts, that are not React specific. Also, help to learn by example.

## Setting up codemod project for TDD workflow

Codemod is textbook sample where TDD works. You have an input file, you run the script and you get output. Thus I would really recommend using TDD for codemod projects. Not only it makes codemods more stable, but having projects with test workflow setup, will help you learn. Because you can experiment just by running the same test over and over again. 

Here is how to create codemod project from scratch:

1. Create empty npm project (`npm init sample-codemod`)
2. Install codeshift `npm i -S jscodeshift`
3. Install jest `npm i -S jest`
4. Copy over test utils from jscodeshift library [src/testUtils.js](https://github.com/facebook/jscodeshift/blob/595b15cda4f52673cad8fe3ae65c12e334966dbf/src/testUtils.js)
5. Modify `testTest.js`, by replacing `require('./core')` with `require('jscodeshift')`
6. Create initial folder structure:

```
+-- src
|   +-- __testfixtures__  - put sample files for transformation, use suffixes .input.js and .output.js
|   +-- __tests__ -simplicity-in-technology.markdown
```

After that, you can create a test file and start adding tests. Test utils from `jscodeshift` allow you to create 2 type tests:

- Inline, where input and output is defined as string `defineInlineTest(transformFn, options, input, output)`
- Using files, where you define path to input and output files `defineTest(__dirname, transformName, options, testFilePrefix)`

I have created a repo with this sample in [Github](https://github.com/katilius/sample-js-codemod).

## Steps to create codemod

Essentially codemods could be oversimplified to just 2 steps:

1. Find the tree node
2. Replace with new one or modify

Since there are many ways to write the same logic in JS. You will need to think of all ways developer could write the thing you want to replace. For example, even finding imported value is not that trivial. You can use `require` instead of `import`, you can rename named import, you can do same import statement multiple times and etc.

At start I would suggest to think only about the simplest case and not think about edge cases. That's why I think TDD is essential, you can gradually add more complex cases, while not breaking initial functionality.

## Sample codemod

Let's write simple codemod using this workflow. First let's define a simple test case, as we are trying to work with TDD.

We want to convert this:

```js
export default (a, b) => a + b;
```

into:

```js
export default function (a, b) {
  return a + b;
}
```

If we are using file approach for jscodeshift. It would be defined this way:

```js
describe('arrow-to-function', () => {
    defineTest(__dirname, 'arrow-to-function', null, 'defaultExportedArrow');
});
```

Once we have this sample, we can launch [AST Explorer](https://astexplorer.net/) and inspect how to input code is parsed as AST (make sure you use esprima spec):

![Arrow AST](https://thepracticaldev.s3.amazonaws.com/i/zm28eoi504t24vic5a34.png)


From explorer it's clear that we need to find the node of type `ArrowFunctionExpression`. And based on the highlight, we care about arrow function `body` and `params` fields.
After analyzing what to find, we also need to find out what we need to build, here AST explorer helps as well. Just paste output code to it:

![Function AST](https://thepracticaldev.s3.amazonaws.com/i/1hqhiz38wur7n5d98lz9.png)

From the structure, it's clear, that regular functions are a bit more complex. We need to add a block statement and return statement.

Let's start with finding arrow functions. To create codeshift transformation you need to create file and export single function. That function will receive three arguments: fileInfo, API, options. Currently, we mostly care about `api.jscodeshift` (usually, it is defined as `j`) and `fileInfo`. Finding all arrow functions is simple:

```js
module.exports = function transform(file, api) {
  const j = api.jscodeshift;

  j(file.source).find(j.ArrowFunctionExpression);
};
```

This will return the collection instance, which we can iterate and replace nodes. Let's replace all arrow functions with regular functions:

```js
module.exports = function transform(file, api) {
  const j = api.jscodeshift;

  return j(file.source)
    .find(j.ArrowFunctionExpression)
    .replaceWith(p => {
      const nodeValue = p.value; // get value from NodePath
      
      // whole node will be replaced with newly built node:
      return j.functionDeclaration(
        j.identifier(""),
        nodeValue.params,
        j.blockStatement([j.returnStatement(nodeValue.body)])
      );
    })
    .toSource();
};
```

- Each item is an instance of `NodePath`, which allows you get parent node, therefore in order to access actual node you need to use `p.value` field.
- If you access jscodeshift field starting with __uppercase__, it will return type (`j.ArrowFunctionExpression`). It is used to filter and check nodes.
- If you access jscodeshift field starting with __lowercase__, it will return build instance. Which allows creating code blocks. Check [AST Types](https://github.com/benjamn/ast-types/) repo to see what fields are supported using each builder. For example, if you would open `core.ts` file and look for `FunctionExpression`, it has following definition: `build("id", "params", "body")`. Which means that you need to pass id, params, and body.

And that's pretty much it. If you follow these steps, it's not that hard to write more complex codemod. Just constantly check __AST Explorer__ and gradually you will become more familiar with the structure.

### Further improvements

The current implementation is extremely naive and __should not be run on actual code base__. Yet, if you would like to work further on this example to learn, here are some suggestions:

- Handle arrow functions with block statement `{}`
- Do not convert arrow functions that call `this`. Arrow functions handle `this` differently and current codemod would break working code.
- Convert arrow function declaration into named functions, for example `const sum = (a, b) => a + b` could be converted to named function `function sum(){...}`


### Running on codebase

I have mentioned previously that this code should not be run on the real codebase, however, if you would build fully working codemod, here is how to run it:

```
npx jscodeshift -t script-path.js pathToFiles
```

### Dealing with complexity

- Extract custom predicates. For example, if you are dealing a lot with JSX, you might create predicates like `hasJsxAttribute`, `isNativeElement`, etc.
- Extract builder functions. If you keep creating import statements, create a function that would return the node with the import statement. 

### Using Typescript

It takes a bit of guessing when using jscodeshift API if you are not familiar with it. Typescript can simplify this process, it works with __AST Types__ mentioned at the start of the post. With Typescript it's a bit easier, to guess what parameters to use in a builder, or how to access certain values. However since parsing is really dynamic in nature, time saved by getting type info is sometimes lost dealing with Typescript type system and defining types manually.

## Jscodeshift recipes  

Here I will share some code snippets, that might help you do some tasks while writing your own codemod. They are not 100% error prone, but at least they show some different modifications you can do. 

### Create function call statement


```js
// will generate this:
const result = sum(2, 2);

j.variableDeclaration('const',
    [j.variableDeclarator(
      j.identifier('t'),
      j.callExpression(j.identifier('result'), [j.literal(2), j.literal(2)])
    )]
  );
```

### Find imports in file

```js
function findImportsByPath(j, root, importPath) {
    const result = {
        defaultImportUsed: false,
        namedImports: []
    };
    root.find(j.ImportDeclaration, (node) => node.source.value === importPath)
        .forEach(nodePath => {
            nodePath.value.specifiers.forEach(specifier => {
                if (j.ImportDefaultSpecifier.check(specifier)) {
                    result.defaultImportUsed = true;
                } else {
                    // specifier interface has both local and imported fields
                    // they are the same unless you rename your import: import {test as b}
                    result.namedImports.push(specifier.imported.name)
                }
            })
        });
    return result;
}
```

### Rename JSX attribute

```js
function transform(file, api) {
    const j = api.jscodeshift;

    return j(file.source)
        .find(j.JSXAttribute, n => n.name.name === 'class')
        .forEach(nodePath => {
            nodePath.node.name = 'className'
        }).toSource();
}
```