Automating the migration of lodash to lodash-es in a large codebase with jscodeshift

Recently the Elements team needed to make a reasonably large change to the codebase: migrating over 300 files which imported lodash to instead import from lodash-es.

To automate this change we chose to write a codemod for jscodeshift, a tool by Facebook. The power of jscodeshift is that it parses your code into an Abstract Syntax Tree (AST) before transforming it, allowing you to write codemods that are smarter than regular expression based codemods.

jscodeshift is a toolkit for running codemods over multiple JS files. It provides:

  • A runner, which executes the provided transform for each file passed to it. It also outputs a summary of how many files have (not) been transformed.
  • A wrapper around recast, providing a different API. Recast is an AST-to-AST transform tool and also tries to preserve the style of original code as much as possible.

Fortunately, we diligently imported Lodash functions from their direct entry points (eg. lodash/<function-name>"), which makes this modification easier.

A typical change we needed to make looks like this:

1
2
3
4
import React from "react"
-import isEmpty from "lodash/isEmpty"
+import { isEmpty } from "lodash-es"
import { connect } from "react-redux"

In many cases, there was more than one import from Lodash to consider.

1
2
3
4
5
import React from "react"
-import isEmpty from "lodash/isEmpty"
-import mapValues from "lodash/mapValues
+import { isEmpty, mapValues } from "lodash-es"
import { connect } from "react-redux"

Writing the jscodeshift transformer

The starting point of any jscodeshift codemod is the transformer function. The transformer function gives you the fileInfo of the file that the CLI is operating on, and the api.jscodeshift API. It’s common to see examples of this API reference aliased to j.

1
2
3
4
5
6
7
8
9
10
11
12
13
export default function transformer(fileInfo, api) {
  // j is a reference to the api we will use later on
  const j = api.jscodeshift

  // Create a jscodeshift Collection from the source string
  const root = j(fileInfo.source)

  // Do some sort of transform on the Collection
  // .. omitted ..

  // Return the new code string
  return root.toSource()
}

When working with the API j, you pass it the file source and it returns a Collection. A Collection is an object containing an array of NodePath objects. The docs describe it as jQuery-like:

jscodeshift is a reference to the wrapper around recast and provides a jQuery-like API to navigate and transform the AST.

Finding the import declarations

The first thing we want to do is find all the import declarations that are sourcing Lodash modules. To better understand how we might do this, it’s helpful to first explore what the AST looks like. The AST for our sample code contains many ImportDeclaration nodes. Each one contains a source string literal with the value, which we will check for Lodash.

The above screenshot was captured in AST Explorer. Check it out, it’s an excellent online tool for exploring ASTs (Turn on Transform -> jscodeshift).

1
2
3
4
5
6
7
8
9
10
const lodashImports = root
  .find(j.ImportDeclaration)
  // Find only lodash NodePaths
  .filter(nodePath => {
    return nodePath.value.source.value.startsWith("lodash")
  })

lodashImports.forEach(nodePath => {
  // ...
})

Analyzing the import declarations

Now that we have all of the Lodash import declarations, we can start to analyze them. We want to understand what the module name is, so let’s further analyze the source literal.

1
2
3
// Eg. source "lodash/mapValues"
const id = nodePath.value.source.value.replace("lodash/", "")
// => "mapValues"

We also want to understand if the import declaration’s specifier is the same as the module name. If it’s not the same, it is important to capture this, because we might need to be able to import the named functions with the as directive in the future.

For example:

1
2
-import myMapValues from "lodash/mapValues"
+import { mapValues as myMapValues } from "lodash-es"

To do this we need to explore what the import specifier AST nodes look like.

1
2
3
4
// Get the first specifier...
const [specifier] = nodePath.value.specifiers
// ...and save the name
const name = specifier ? specifier.local.name : id

With that, we now have enough information to create new import specifiers. As we loop over each import declaration, we’ll populate an array of replacement specifiers we intend to use on our final import declaration.

1
2
3
4
5
const replacementSpecifier = j.importSpecifier(
  j.identifier(id), // eg. isEmpty or mapValues
  j.identifier(name) // eg. isEmpty or myMapValues
)
replacementSpecifiers.push(replacementSpecifier)

Removing and replacing import declarations

Now that we have all of the replacement specifiers, we can look at replacing all the lodash/* import declarations we found, with a single import declaration for lodash-es. To do so we’ll need maintain a reference to the first import declaration for later. All the other Lodash import nodes can be removed.

To remove nodes in jscodeshift, we’ll need to wrap it in a j(nodePath) call and use remove().

1
2
3
4
5
if (!first) {
  first = nodePath
} else {
  j(nodePath).remove()
}

Using that first Lodash import reference, we can create the new import declaration.

To replace nodes in jscodeshift, we’ll need to wrap it in a j(nodePath) call and use replaceWith()

1
2
const newImport = j.importDeclaration(replacementSpecifiers, j.literal("lodash-es"))
j(first).replaceWith(newImport)

Alternatively you can modify the the first Lodash import reference itself. This is my preferred way to make this change.

1
2
first.value.specifiers = replacementSpecifiers
first.value.source.value = "lodash-es"

Dealing with comments

The modifications we’ve applied so far are, are enough for our application code to be functional, but we will lose all the comments from the nodes with removed and replaced since, in the AST, comments are attached to nodes. In order to fix this problem, we’ll need to collect up all the comments and assign them to the final lodash-es import declaration.

1
2
3
4
5
6
7
let replacementComments = []
// Save all the comments
lodashImports.forEach(nodePath => {
  replacementComments = comments.concat(nodePath.value.comments || [])
})
// Replace the comments
first.value.comments = replacementComments

Putting it all together

That’s it, now lets assemble it into our final transform function.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
export default function transformer(fileInfo, api) {
  const j = api.jscodeshift
  const root = j(fileInfo.source)

  const lodashImports = root
    .find(j.ImportDeclaration) //
    .filter(nodePath => {
      return nodePath.value.source.value.startsWith("lodash")
    })

  let first
  let replacementComments = []
  const replacementSpecifiers = []

  lodashImports.forEach(nodePath => {
    const id = nodePath.value.source.value.replace("lodash/", "")
    const [specifier] = nodePath.value.specifiers
    const name = specifier ? specifier.local.name : id
    const replacementSpecifier = j.importSpecifier(
      j.identifier(id), // the import id
      j.identifier(name) // the import "as" name, it might be the same as id.
    )
    replacementSpecifiers.push(replacementSpecifier)
    replacementComments = replacementComments.concat(nodePath.value.comments || [])
    if (!first) {
      first = nodePath
    } else {
      j(nodePath).remove()
    }
  })

  if (first) {
    first.value.specifiers = replacementSpecifiers
    first.value.source.value = "lodash-es"
    first.value.comments = replacementComments
  }

  return root.toSource()
}

Finally, run the codemod over your entire codebase with the CLI:

1
jscodeshift src -t lodash-es-imports.js --extensions=js,jsx

And the result? Checking the git diff, we see many changes like this:

1
2
3
4
5
6
7
import React from "react"
// isEmpty used because it works on arrays and objects
-import isEmpty from "lodash/isEmpty"
// Note that myMapValues is different from mapValues
-import myMapValues from "lodash/mapValues"
+import { isEmpty, mapValues as myMapValues } from "lodash-es";
import { connect } from "react-redux"

You can copy the complete example above and play with it in AST Explorer. You’ll need to turn on Transform -> jscodeshift in order to see the codemod output.

Happy codemod’n.

Links