Writing codemods with jscodeshift

John Daly

6/4/2022

When working with large software projects, you will inevitably encounter situations where you need make updates that affect many pieces of your codebase. A common scenario is updating code that uses a deprecated API. You might also want to adopt a new language feature or pattern that improves the readability and/or performance of your code.

To complete these code migrations, you will often go in an make the necessary changes manually, which can be time consuming and tedious. Fortunately, there is a way to accomplish these types of migrations with code: codemods.

The idea behind codemods is to define a set of rules that dictate how to transform source code. Write the rules once, and then apply to all relevant files in your project.

In the JavaScript ecosystem, there is a great library called jscodeshift that provides an API for writing codemods. In this post, I'll share an anecdote of how I used codemods to solve a problem at work, and I'll break down how to write that codemod using jscodeshift.

Setting the stage

I'm part of the Web Platform team at Convoy, and I work on tools that help teams build web applications. One of the codebases that my team maintains is the design system library.

One of the product teams that I work with had noticed that development compile times had increased significantly when updating to a newer version of the design system. After some investigation, we were able to narrow the problem down to a single component, the Icon.

Here's a simplified version of what the component looked like. Can you spot the problem?

import * as Icons from '@material-ui/icons';
import * as Styled from './Styled';
type Props = {
  iconName: keyof typeof Icons;
  className?: string;
}
const Icon = ({ iconName, className }: Props) => {
  const MaterialIcon = Icons[iconName];
  return (
    <Styled.IconContainer className={className}>
      <MaterialIcon />
    </Styled.IconContainer>
  );
}

The issue is the import * as Icons from '@material-ui/icons' statement. Its causing all of the icon components from @material-ui to be included in the development bundle. This is a problem, because Material UI has a LOT of icon components. At the time of writing, Material UI has 5000+ icon components. That's 5000+ components that Webpack has to process when creating a development bundle. Each icon component module is small, so its really a "death by 5000 cuts" scenario. The Material UI docs even have a guide, with suggestions for how to mitigate the problem.

To demonstrate the impact, using my personal website as an example app, I added the following code to the main page:

import * as Icons from '@material-ui/icons';
const IconByName = (iconName: keyof typeof Icons) => {
  return Icons[iconName];
}

Next, I ran the @next/bundle-analyzer Webpack plugin on my development build, and got the following result:

Webpack includes all of the icon components from @material-ui/icons

Even if I'm using a single icon from @material-ui/icons, the development bundle includes all 5000+ icons in it. 😱

Unfortunately, it gets even worse. Since we are using the Icons[iconName] pattern, Webpack is unable to tree-shake unused icons from our production bundle, and we end up with 5000+ icons in our production bundle. Ouch.

We need to solve the production issue first, since it negatively affects an end user's experience. We'll need to update our code so that the set of icons we need can be determined through static analysis at build time.

A possible solution would be to update our import statement to explicitly state the icons we want:

import { Info } from '@material-ui/icons';

This will work for production, but not in development. It won't work for development, because Webpack will load the module at @material-ui/icons, which is a file that re-exports all of the individual icons. Webpack won't perform tree-shaking optimizations in development, so it will still include all of the icons in the development bundle. We'll have to use a different approach.

Using path imports will address the problem for both development and production builds:

import Info from '@material-ui/icons/Info';

This made for a dramatic improvement to the size of the development bundle:

Bundle size with individual Material UI icon imports

Webpack only includes the icon components I have specified in the bundle now!

Fixing our Icon component's API

With a solution to the underlying problem identified, we needed to update our approach to the design system's Icon API.

Instead of importing the Material UI icons in the design system library, we would have the consumers of the library provide the icons to the component. This allows product teams to have control over which icons they are importing into their bundles.

ℹ️ You can see what the API change would look like, with this interactive example

import { Icon } from '@convoy/fuel';
const ComponentWithIcon = () => {
  return (
    <Icon
      iconName={'ArrowRightSharp'}
    />
  )
}

Before

Passing in the icon name as a string, through the iconName prop

After testing this change, we found that the start up times were significantly faster 🙌.

	Before (All Icons in Bundle)	After (Individual Icons in Bundle)
Modules compiled (dev):	5905	353
Compile time (dev):	11.2s	3.5s
Total bundle size (dev):	5.52 MB	2.62 MB

A new problem

Making these changes introduces with a new problem: This was going to be a breaking API change, and every app that used the Icon component would need to use the updated code.

We couldn't keep the old approach in, for backwards compatibility, since it would still result in apps importing all of the icons in development, even if they were using the new API.

Icon usage was pretty heavy in lots of apps, with some apps even having hundreds of instances 😱. On top of that, the migration wouldn't be as simple as just using find/replace, since we would need to add the import for the icon component from Material UI.

Enter codemods and jscodeshift!

This migration was a very good use case for learning how to write codemods using jscodeshift. We could define the codemods once, and then apply them to all of the applications that needed to be updated.

The key with codemods is to identify a common repeatable pattern, which you'll apply to all files in the project.

ℹ️ Use the controls below to see the steps that our codemod will take to transform the code

import React from 'react';
import Fuel from '@convoy/fuel';
const Example = () => {
  return (
    <Fuel.Icon
      iconName={'ArrowRightSharp'}
    />
  )
}
export default Example;

Step 1

Open the file, so that we can find the code we'd like to transform

Building our codemod

Now that we have our general idea in place, we can start to write the codemod!

ℹ️ This interactive example will take us step by step through the process

import type { Transform } from 'jscodeshift';
// This codemod uses the 'tsx' parser
export const parser = 'tsx';
const transformer: Transform = (file, api, options) => {
}
export default transformer;

The basis of a jscodeshift codemod is a Transform function, where we define the transformations that we will be performing on the source code

Here's the entire implementation, with added comments:

import type { Transform } from 'jscodeshift';
// This codemod uses the 'tsx' parser
export const parser = 'tsx';
/**
 * Will update usage of Icons to use the preferred path import syntax
 */
const transformer: Transform = (file, api, options) => {
  const j = api.jscodeshift;
  const printOptions = options.printOptions || {
    quote: 'single',
    trailingComma: true,
  };
  const root = j(file.source);
  // Keep track of the icons that are being used
  const iconsUsed = new Set<string>();
  let wasAnIconFound = false;
  // Find what the Fuel import is called
  let fuelAlias = 'Fuel';
  root
    .find(j.ImportDeclaration)
    .filter(path => path.value.source.value === '@convoy/fuel')
    .find(j.ImportNamespaceSpecifier)
    .filter(path => path.value.local.type === 'Identifier')
    .forEach(path => {
      if (path.value.local.name) {
        fuelAlias = path.value.local.name;
      }
    });
  root
    .find(j.JSXElement)
    .filter(path => {
      // <Fuel.Icon />
      // <FuelAlias.Icon />
      const isFuelIcon =
        path.value.openingElement.name.type === 'JSXMemberExpression' &&
        path.value.openingElement.name.object.type === 'JSXIdentifier' &&
        path.value.openingElement.name.object.name === fuelAlias &&
        path.value.openingElement.name.property.name === 'Icon';
      return isFuelIcon || isFuelButtonIcon;
    })
    .forEach(path => {
      path.value.openingElement.attributes.forEach(attr => {
        if (attr.type === 'JSXAttribute' && attr.name.name === 'iconName') {
          // Update 'iconName' prop to be 'icon'
          attr.name.name = 'icon';
          // Input:   <Fuel.Icon iconName={'ArrowUpwardSharp'} />
          // Output:  <Fuel.Icon icon={ArrowUpwardSharp} />
          if (
            attr.value.type === 'JSXExpressionContainer' &&
            attr.value.expression.type === 'StringLiteral'
          ) {
            const iconName = attr.value.expression.value;
            iconsUsed.add(iconName);
            attr.value.expression = j.identifier(iconName);
          }
          // Input:   <Fuel.Icon iconName='ArrowUpwardSharp' />
          // Output:  <Fuel.Icon icon={ArrowUpwardSharp} />
          else if (attr.value.type === 'StringLiteral') {
            const iconName = attr.value.value;
            iconsUsed.add(iconName);
            attr.value = j.jsxExpressionContainer(j.identifier(iconName));
          }
        }
      });
    });
  // Find the original import, which will mark where we insert the new imports
  const originalImport = root
    .find(j.ImportDeclaration)
    .filter(path => path.value.source.value === '@convoy/fuel')
    .at(0);
  // Loop through the list of Icons and add the `@material-ui` import for each
  const iconsUsedList = Array.from(iconsUsed);
  for (const icon of iconsUsedList) {
    originalImport.insertAfter(
      j.importDeclaration(
        [j.importDefaultSpecifier(j.identifier(icon))],
        j.stringLiteral(`@material-ui/icons/${icon}`),
      ),
    );
  }
  // If we found any icons, then we want to commit the transforms to the source file
  if (iconsUsedList.length > 0) {
    wasAnIconFound = true;
  }
  return wasAnIconFound ? root.toSource(printOptions) : null;
};
export default transformer;

With this codemod, we were able to update all the applications that used the design system. A couple hours to write the codemod saved days of manual migration work!

If you want to get started learning more about codemods, I highly recommend checking out AST Explorer. It allows you to: inspect the AST of your source code, write codemods, and apply those codemods to source code, all in the browser! Its a great tool for learning how ASTs are structured, and for prototyping codemods. There is an initial learning curve, but getting the basics down could save you a lot of time in your future code migration work.