Writing a custom Babel plugin using learnings from React

When working on a component library, we sometimes need to automatically and consistently include certain features in every component. The goal: make working with the library easier for developers.

Photo by Suzanne D. Williams on Unsplash

One common use case my design system dev team at Nasdaq recently encountered was to ensure that we:

• Provide effective support to the more than a dozen applications consuming our library when they use components as intended
• Warn our users that custom style overrides on our components would effectively “void the warranty” — they would need to proceed at their own risk and could potentially break intended styles by applying their own

Possible approaches

Arguably the simplest solution could be to add a warning message to the library readme. The problem here is that this relies on every developer being proactive in reading the readme and retaining this bit of information for when it becomes relevant later on, a tall order when they may be working with many libraries beyond our own.

Another approach could be to create a React hook and require that it be used in every component that defines its own styles. This time, we move to the other extreme of maintainability, adding substantial overhead to virtually every component we create, making scaling the library much more cumbersome.

Yet a third option may be to create a higher-order component (HOC) and mandate that each component be wrapped with it before being exported. While somewhat less cumbersome than using hooks, this approach also has the drawback of needing to be implemented directly in every component.

Given the drawbacks of the various options outlined above, let’s take a look at a fourth approach: writing a custom Babel plugin.

Babel and ASTs to the rescue!

As described on their website, Babel is a JavaScript compiler that takes in what is typically a newer version of JS and compiles it down to a syntax that most browsers can understand natively. I won’t get into a full discussion of how Babel works here, but if you’re not familiar with the library, have a look at these resources to get started.

While the primary use case for Babel is compiling JavaScript to an older syntax, another important feature of the library is giving developers the ability to write plugins to transform code using abstract syntax trees (ASTs). Again, an introduction to the concept of ASTs is beyond the scope of this article, so here are some resources to get you started.

In what follows, we will create our own Babel plugin to automatically wrap components with a higher-order React component at the build stage.

We needed to ensure that we warn our users that custom style overrides on our components would effectively “void the warranty”

Warning users when they use “dangerous” props

Right off the bat, we knew that the two primary methods that developers would use to override our built-in styles would be:

• Using the style prop to define inline style overrides
• Using the className prop to refer to an additional CSS class

Of course, there are other ways to implement styles, such as using the id prop or even custom HTML attributes. However, we estimated that for an initial implementation, the style and className props would cover >80% of the use cases that we wanted to address.

Having nailed down our initial focus, we decided that developers should get console warnings whenever they attempted to use these props. Additionally, to keep these messages from becoming too intrusive, we defined the following acceptance criteria:

• Style-related console warnings should only be shown in development environments
• Users should be able to silence all style-related console warnings globally, if desired
• Users should be able to silence style-related console warnings for individual components, if desired

Now lets have a look at how we might implement this solution.

Creating a higher-order component (HOC)

To enable consistent messaging everywhere, we decided that our approach would be to create a HOC and wrap all default exports in files containing a styled-component import with it.

Additionally, to satisfy the acceptance criteria outlined above, we also added the following features:

• Only create style-related console warnings when process.env.NODE_ENV !== "production" and on component mount, or if one of the relevant props changes
• Silence all style-related console warnings from every affected component if REACT_APP_SILENCE_ALL_STYLE_ALERTS is true
• If a specific component is given the accept-custom-class-or-style as a className, then silence alerts for that component

The resulting HOC looks like the following:

https://gist.github.com/haldunanil/170e4e05b28157c56eb0771589380a2d#file-withstylealerts-tsx

To make this component easily accessible everywhere, we will publish it to NPM and add it as a dependency for our project. If you’re not familiar with how the process of publishing a new package to NPM works, have a look here and here.

Any component wrapped with withStyleAlerts will now get the desired features above! Our next step is to create a Babel plugin to automatically wrap all our styled component exports with his new HOC.

Writing our Babel plugin

Before we get into our code, let’s once again go through our objectives for what we want it to do. In order for an exported component to be wrapped with our HOC, it needs to satisfy the following criteria:

• A default import of a CSS-in-JS library is used ( styled-components by default)
• There is a default export in the file (this should correspond to the component being created)

Additionally, we would like the name of our HOC and the library it’s exported from to be defined by the user. Once all of these conditions are satisfied, we expect a file like this:

import styled from "styled-components";

const MyComponent = styled.div`
  position: relative;
  width: 100%;
  height: 4px;
  border-radius: 0.5rem;
`;

export default MyComponent;

To be transformed into this (changes in bold):

import styled from "styled-components";
import { withStyleAlerts } from "test-lib-name";

const MyComponent = styled.div`
  position: relative;
  width: 100%;
  height: 4px;
  border-radius: 0.5rem;
`;

export default withStyleAlerts(MyComponent);

Finally, we need to think about how we’re going to traverse our component files. We have to check multiple different places in a component file as preconditions of whether we want to actually make changes to it; for instance, a default import of the desired CSS-in-JS library is needed in order to wrap the component with our HOC.

With that out of the way, let’s start writing our plugin!

Step 0a: Setting up the base plugin structure

First, we want to ensure that the name of the HOC (e.g. withStyleAlerts ) and the name of the library that exports it (e.g. test-lib-name ) are defined as plugin options. If not, the plugin should immediately return. To implement this, we can do:

Our initial conditional to prevent unnecessary operations

In order to use this plugin (which is named babel-plugin-wrap-components-with-hoc, but you can name it whatever you’d like), we need to edit our .babelrc file (or equivalent) and add the following line to our plugins array:

Add the plugin to the plugins array in .babelrc

Step 0b: Borrowing state from React

As mentioned earlier, we need to keep track of a number of things when making changes to our component files. These include:

• Whether the HOC and/or another export from its library already exists in the file, which itself could result in three possible scenarios:
  1. HOC is already imported, no change necessary
  2. HOC is not imported but another export from its library is, add to import array
  3. Nothing is imported from the HOC library, add import statement
• Whether a default export exists in the file
• Whether the designated CSS-in-JS library is default imported

To keep track of all of this, we want to create a state for each file that is being traversed by Babel. Just like React, the state should be high-enough that it’s accessible as the single source of truth for all our traversals. Noting that having an external state as shown here in the Babel plugin handbook is likely to cause unexpected behavior and hard-to-diagnose errors, we will instead be maintaining state in our Program visitor as shown here:

Adding state to our Program node

We’ll see how each variable in our state will be used shortly.

Step 0c: Borrowing state change handlers from React

AST node traversals in Babel happen top down, meaning that it’s not possible to do bidirectional data transfers. However, we still need to be able to update our state from inside a node further down the tree.

To update our state from child nodes, we’ll once again take a page out of the React playbook and create state change handlers in our Program node. Subsequently, we’ll pass these down similar to how we pass down component methods using props. More on passing them down in a bit, let’s first create our change handlers:

Adding state change handlers to the Program node

Let’s now turn our attention to traversing our nodes and updating our file as desired.

Step 1a: Traverse import declarations

Let’s begin our traversal by checking every ImportDeclaration node, which corresponds to the an entire import statement like import styled from “styled-components";for two things:

1. Whether the HOC library is imported anywhere in our file
2. Whether the CSS-in-JS library is default imported in our file

If we find that the HOC library is being imported, we’ll also note the path of that import and set it as our styleUtilsPath. If we find that the CSS-in-JS file is default imported, then we set the setDefaultImportAvailable to true. Our resulting code then looks like something like this:

Adding first traversal to get ImportDeclaration nodes and update state

Once again, it’s important here to highlight that the ImportDeclaration node wasn’t added inside the visitor object alongside Program , but was called using the path.traverse function inside Program. The reason we did this is to enable the sharing of a collective state within Program, as this makes our steps synchronous and therefore predictable.

Note that the second of the argument passed into path.traverse is an object containing state and two of our state change functions. Personally, I find that it helps to think of this object as the props of the GetImportDeclarationNode object. Not to potentially beat a dead horse, but seeing the similarities between these two functions is actually pretty cool:

# babel traversal function
path.traverse(GetImportDeclarationNode, {
  state,
  setStyleUtilsPath,
  setDefaultImportAvailable
});
# example React component using `createElement` function
React.createElement("div", {
  isOpen,
  style,
  className
});
# same React component using JSX
<div isOpen={isOpen} style={style} className={className} />

These days, most of us typically use the JSX syntax shown in the last example above to create a React component. However, what the browser sees after we build our React apps and Babel comes into the picture is actually the second example using React.createElement.

I hope this helps you establish a mental model around how we can use learnings in React to think about traversing AST nodes in Babel. Who knows, maybe we’ll even start using something like JSX when traversing Babel nodes! I’d call it BabelX. 😜

Step 1b: If CSS-in-JS library not imported, move to next file

Since we’re using the presence of the CSS-in-JS library in a file to indicate that it contains a styled component (some pun intended!), we won’t apply any transformations if we find that the library is not used (i.e. isDefaultImportAvailable is false). With this addition, our file becomes:

Skip file parsing when CSS-in-JS default import not found

Step 1c: Check if our HOC library is imported

It’s time to begin making changes to our file. Let’s begin by importing the withStyleAlerts HOC when it is not present. To avoid having multiple imports from test-lib-name, we’ll only execute this step if our library is not mentioned in our file at all. There will be additional steps later on to take care of the scenarios where test-lib-name is mentioned (e.g. if you’re using a single library as a styling utility resource).

We can accomplish these by making the following changes to our file:

Adding HOC to top of file if library not currently referenced

Note that we’re using the path.unshiftContainer function to add a new import declaration to the top of our file. Because we will be referencing the withStyleAlerts name at the end of our file when we call it with our default exported component, we also need to create the appropriate identifier and import specifier (see newly added step 0d).

Step 2: Take care of HOC import edge cases

Now, let’s ensure that we’re not duplicating any imports and adding them where necessary:

Adding HOC import to existing imports from its library, if they do exist

Note here that two two sub-steps in 2 accomplish the following:

• Step 2a: Check if withStyleAlerts exists as an identifier among the imports and set the state to reflect the result
• Step 2b: If withStyleAlerts still does not exist as an identifier among the imports, then push it to the list of imports from the library containing it

Thanks to the combined efforts of steps 1a-1c and 2a-2b, we have now covered all scenarios surrounding the import of our HOC and ensured that it is correctly imported by the time we reach step 3.

Step 3: Wrap default export with HOC

The final step of our Babel journey consists of wrapping our default exported component with our withStyleAlerts HOC. Before we get into the specifics of our implementation, however, I want to quickly remark on a caveat with this approach.

The solution as presented here assumes that all components that should be wrapped will be default exported. The reason for this is twofold:

• We want to encourage the modularization of individual components, for ease of testing as well as simplification of maintenance overhead
• Dealing with named exports is much trickier, since they can include a host of other objects that are exported for reuse but aren’t actually React components

Given these caveats above, note that the following implementation may not work as desired if you typically export many components from the same file, do not use default exports, or use default exports in component files for something other than exporting components.

With that out of the way, let’s look at our finished implementation:

https://gist.github.com/haldunanil/170e4e05b28157c56eb0771589380a2d#file-plugin-js

Let’s break down steps 3a and 3b. The primary objective of 3a is to traverse the AST tree and find out export default path. This will allow us to replace the right side of the export with a function call that passes our current default exported component as an argument of withStyleAlerts.

Step 3b actually implements this replacement, by constructing the call expression and then setting the export default’s declaration equal to it.

That wraps up the implementation of our Babel plugin! To use, you’ll simply need to deploy your plugin and a library containing the withStyleAlerts HOC to NPM and you’ll be all set to begin transforming your files.

A word of caution

Using Babel plugins to do custom transformations beyond syntax changes as shown above is a powerful way to modify your code. However, as Uncle Ben put it many times in many different Spider-Man movies: with great power, comes great responsibility.

Photo by Bruno Nascimento on Unsplash

Changes applied during the build process will likely be missed by most testing suites which typically use the uncompiled source code. Using Babel plugins to transform your code and add extensive functionality could result in bugs that are hard to diagnose and address, as they will often only appear once your library is being used elsewhere.

Additionally, developers working on a library that uses Babel plugins to transform code may not be aware that such a plugin is even in place — in fact, it’s probably better to assume that they won’t be aware. Again, this might result in unexpected behavior that may be hard to pin down, potentially leading to confusion and frustration among your developers.

As a general rule of thumb, it’s good to try to stay away from using Babel plugins to make destructive changes or modifications to existing features. It’s also a good idea to ensure that your new functionality is well encapsulated and is designed in a way that is not likely to overlap with existing and future library functionality.

Conclusion

I hope this article was able to demonstrate that Babel isn’t just for transforming newer JavaScript syntax to an older one that browsers can work with, but can also be used to simplify development processes and include a ton of features automatically during your build process.

I also hope that this article shows how a working knowledge of React can help when creating Babel plugins and that doing so is not as scary a prospect as it may initially seem!

Feel free to leave comments below, get in touch with my via Twitter, and check out my website!

The full gist containing the HOC and the plugin above can be found here.