JS Refactors at Scale? Meet Codemods + AST

We have all been there when we have to upgrade a library and there are lot of changes we need to do manually to each file. therefore some libraries provide codemods (like a script) that automatically updates the deprecated code required for the new version of that library. I know we can achieve this by node file system, but there is an issue where we have to specify directly what we want to change or copying the arguments of a function, but codemods are superfast as they run on multiple threads. Enough with the chit chat and let’s build one for ourself.

What is a “Codemod”?

By doing a google search, Codemods are transformations that run on your codebase programmatically. This allows a large number of changes to be programmatically applied without having to manually go through every file.

Let’s simplify it. if we break the word codemod it kind of sounds like “code modification” or in short codemod. we are creating a script that goes through our file and modifies our code. It’s as simple as that.But you can thing why not regex, we can also target via regex and write a node script that does the same thing. Regex works great for text but not for code. Regex doesn’t understand structure. for e.g when we are updating some random component like say a Button

<Button primary>Click Me!</Button>

<Button variant="primary">Click Me!</Button>

and you want to transform Button primary attribute to variant=”Primary” . With the regex approach it will look something like this code.replace(/<Button\s+primary(.*?)>/g, '<Button variant="primary"$1>') , Now imagine you have to target <Button secondary>Click Me!</Button> or success.It becomes very hard from regex to manage all this scenario. Now to handle these situations we are going to need to use AST(Abstract syntax tree). another jargon huh. Let’s understand AST

What is AST?

According to google `An Abstract Syntax Tree (AST) is a tree representation of the abstract syntactic structure of a code snippet or formal language. A lot of jargon and a little hard to understand. Lt’s simplify it. Think of an AST is like an x-ray of your code. It doesn’t care about how your code looks but focuses on what it’s mean.

For e.g const sum = a + b take this code now how it will look as an abstract tree

VariableDeclaration (const)
├── VariableDeclarator
│   ├── Identifier (sum)
│   └── BinaryExpression (+)
│       ├── Identifier (a)
│       └── Identifier (b)

This is the AST.

so instead of working with raw text (which is messy and error-prone), you get a clean, structured representation that’s much easier to traverse, analyze and modify safely.

If you want to visualize your code to AST you can go to this website link.

Codemod Tools

There are some popular codemod tools that helps us streamline the process.

• jscodeShift (by facebook): https://github.com/facebook/jscodeshift

• ts-morph (for TypeScript-heavy codebases) : https://ts-morph.com/

• babel plugins (for custom transformations) : https://babeljs.io/

• recast and ast-types (used under the hood)

For the purpose of this blog we’ll just focus on jscodeshift.

Writing your first codemod

Let’s start with the simple like renaming a variable. Here I’ll take the example of upgrading react-router-dom from v5 top v6

In react-router-dom v6 there is no useHistory and we have to rename it to useNavigate.

First let’s install jscodeShift in our project where we want to change our code. Also you can go through the jscodeShift docs here https://jscodeshift.com.

npm install jscodeshift

Second create a javascript file in the root name migrate.js(you can name anything) and inside that create a module that exports a function and accepts two parameters file and api.

module.exports = function transformer(file, api) {
}

Now we need to initialize the jscodeshift and file source inside the transformer function to convert the source code to AST.

module.exports = function transformer(file, api) {
    const j = api.jscodeshift;
    const root = j(file.source);
}

let’s see what going on

• const j = api.jscodeshift; we are assigning the jscodeshift library to the variable j as an utility function.We’ll use this for parsing, traversing and transforming js code as an AST.

• const root = j(file.source); It will parse the source code of the file into an AST and wraps it in jscodeshift collection. Now root is the entry point for finding and transforming nodes in the code.

now we want to change the const history = useHistory() → const navigate = useNavigate(), now to do so we’ll visualize using this link https://astexplorer.net/

Now this is how this simple react program will look like.

Let’s deep dive the right part. we can see we have two ImportDeclaration as in the left we have two import for react and react-router-dom. Then there is a VariableDeclaration and that is CustomComponent

const CustomComponent = (props)=>{
  const history = useHistory();
  const location = useLocation()

  return <button onClick={(e) => history.push(`/hello/${e}`)}>click me</button>
}

now if we expand the VariableDeclarator we get a lot of info, but if you can see the highlighted part we have two more VariableDeclaration and a ReturnStatement.Now if we see on the left part of the code, we’ll see

const history = useHistory();
const location = useLocation()

return <button onClick={(e) => history.push(`/hello/${e}`)}>click me</button>

these history, location and the return statements are the VariableDeclaration and the ReturnStatement.

Now we need to find the node that contains the history and useHistory and to do that jscodehift provides us a method called find(type) it takes what type of node you want to finds. Since, we wants VariableDeclaration will pass this as a type.

module.exports = function transformer(file, api) {
    const j = api.jscodeshift;
    const root = j(file.source);

    const variableDeclarations = root.find(j.VariableDeclaration);
}

Now the variableDeclarations is a an array as it will have all the variableDeclarations node, so we need to loop through it and reach to the correct node or property

Now looping thorugh variableDeclarations will give a Collection of nodepaths which has a .value property.Also Declarations is also an array So we need to loop through again to go to the history and that is inside declarations[i] > id > name

module.exports = function transformer(file, api) {
    const j = api.jscodeshift;
    const root = j(file.source);

    const variableDeclarations = root.find(j.VariableDeclaration);

    variableDeclarations.forEach((declaration) => {
        declaration.value.declarations.forEach((declarator) => {
              if (declarator.id.name === 'history') {
                    declarator.id.name = 'navigate';
              }
        });
  });
}

similarly for the useHistory it is inside the declaration > init > callee > name

so now to replace that we need to update the name.

module.exports = function transformer(file, api) {
    const j = api.jscodeshift;
    const root = j(file.source);

    const variableDeclarations = root.find(j.VariableDeclaration);

    variableDeclarations.forEach((declaration) => {
        declaration.value.declarations.forEach((declarator) => {
              if (declarator.id.name === 'history') {
                    declarator.id.name = 'navigate';
                    declarator.init.callee.name = 'useNavigate';
              }
        });
  });
}

It’s almost done. we have completed the process of renaming the history → navigate and useHistory → useNavigate. The only then left for us is to return the output from AST back to original source code(js)

and to do that we need to return root.toSource(), where root is the transformed AST and toSource() is converting the AST back to js

module.exports = function transformer(file, api) {
    const j = api.jscodeshift;
    const root = j(file.source);

    const variableDeclarations = root.find(j.VariableDeclaration);

    variableDeclarations.forEach((declaration) => {
        declaration.value.declarations.forEach((declarator) => {
              if (declarator.id.name === 'history') {
                    declarator.id.name = 'navigate';
                    declarator.init.callee.name = 'useNavigate';
              }
        });
  });

  return root.toSource()
}

Now let’s create a script in package.json to run this codemod, the command will look like this

jscodeshift -t <jsocdeshit-script-path-we-created> <source-of-file-we-want-to-update>

  "scripts": {
    "migrate": "jscodeshift -t migrate.js index.js"
  }

or if we want to pass the source file we can make like this,

  "scripts": {
    "migrate": "jscodeshift -t migrate.js --"
  }

npm run migrate src/pages/home

let’s see the magic

Now we can even make the import statement to change from the useHistory → useNavigate.

Hint: use find(j.ImportDeclaration).

Conclusion

Codemods and AST provide a powerful tool to automate code refactoring at scale. Something to keep in mind is that have some kind of versioning system like git to help some mishaps. Also starts with small repo before running on /src.

Hope you have learned something and will build something cool for your usecase.

Drop a comment, or ping me on twitter/linkedIn. I’d love to see what you are building.