How to Mutate an AST and Automatically Replace Code Components

[00:00 - 00:12] The script from the previous module, red files from the codebase, traverse the source code within those files to find the relevant code, or nodes in an AST, and log that data. With the data, we were able to define and build a component.

[00:13 - 00:20] Now that we have the component, it's time to go through and replace all of the button elements with our new button component. There are a few ways this could be approached.

[00:21 - 00:31] First, manually, we could go through all the files, look for button elements, and replace them with the button component. This is good for complex changes, but only works well in smaller codebases.

[00:32 - 00:39] Second, find and replace. Use a tool that can find, then replace pieces of code based on text or a regular expression.

[00:40 - 00:45] This is reasonable for simple cases, and works well in large codebases. And finally, ASTs.

[00:46 - 00:52] This is ideal for larger codebases, and can also handle complex transformations . All of these are valid approaches.

[00:53 - 00:59] The scope and complexity of the change can help determine which is the most efficient. The button transformation needed here is complex.

[01:00 - 01:13] For example, the class name string needs to be broken up into multiple props, but those props are only required if they're not the same as the default. Since the sample codebase only contains a handful of buttons, it would be realistic to manually replace them all.

[01:14 - 01:25] However, we're going to use AST-based tooling to understand how it can be applied to transform code. Switching back to the terminal, let's again get started by creating a new sibling directory to the flashcards app.

[01:26 - 01:33] Then we can change into that new directory. And then initialize a new package.json file.

[01:34 - 01:52] Let's start by adding the same dependencies as the previous lesson along with the same types. We're also going to go ahead and add two more Babel packages which we'll cover in more depth later in the lesson.

[01:53 - 02:01] These are the Babel types package in the Babel's generator package. This will help with creating new nodes and converting an AST back into source code.

[02:02 - 02:17] We also need to install the types for the Babel generator package. Then we'll copy all of the logic from the audit.ts script from the previous lesson, but we'll remove all of the logging.

[02:18 - 02:29] This logging was only necessary for performing the initial audit. We now have the basics of our transform script.

[02:30 - 02:44] Much of the logic is identical to the previous audit.ts script because we want to target the identical nodes to transform. Currently, the node's name, node.name.name is button, but now we need to change it to our uppercase button component.

[02:45 - 02:59] The first step is to reassign this name. Now that we've updated the node's name to the button component and mutated the AST, the tree can be turned back into source code using the Babel generator package 's generate method.

[03:00 - 03:12] This accepts an AST as the input and will return a string of code that is equivalent to the AST. Finally, we can overwrite the existing file with this new code string.

[03:13 - 03:29] Running this should now make a transformation in the Flash app codebase. Now let's open a file that contained a button element that also had the button class name to verify that the transformation ran.

[03:30 - 03:44] For example, if we look at the file in the source component's logging directory named log_in.js, the opening element of one of the buttons is now an uppercase button . After running this transform script, we've just created our first code mod.

[03:45 - 03:58] The script converted the code into an AST, changed the AST, and then converted the mutated AST back into code. However, if you look at some of the transformations, there are still a number of issues.

[03:59 - 04:12] First, only the opening element was renamed, so the closing element is still lowercase button. Second, the props are no longer correct since class name is no longer an option , and many have defaults that make them no longer required.

[04:13 - 04:23] Recall from implementing the button component that using the variant in block props instead was done for ease of maintenance and styling consistency. Third, the button component isn't imported.

[04:24 - 04:31] And fourth, some of the formatting changed and some of the files after the transformation was ran. Let's tackle each of these problems individually.

[04:32 - 04:39] The first three are requirements for this transformation to run successfully. The last is optional since the code is still correct, but formatting is lost.

[04:40 - 04:52] Before we continue with updating the script, let's take a moment to copy this transformed code and paste it into AST Explorer to explore the tree. Pause the video and take a moment to see if you can notice what the problem is.

[04:53 - 05:02] Pasting the code as is will result in an error since the opening closing elements don't match. Let's update the closing element to match since this is the desired end state.

[05:03 - 05:14] Using the tree, we can see that we mutated the name property on the name node of the opening element node. The opening element is a property on a JSX element node.

[05:15 - 05:27] The JSX element node also has a closing element property. Inspecting this node, we see that we also need to update this name property on this node's name node.

[05:28 - 05:36] Our traversal script is currently targeting JSX opening element nodes. This means it's not easy to directly reference the closing element node.

[05:37 - 05:57] It is possible to reference parent nodes from the path object, but this would require both traversing up and down the tree, and it's usually easiest to target the highest node in the tree while also trying to be as specific or low as possible in the tree. A more technical definition would be to target the closest shared ancestor of the two nodes.

[05:58 - 06:06] For example, the script needs to modify both the opening and closing elements. They are both properties on the JSX element node.

[06:07 - 06:17] In other words, they are both child nodes of the JSX element node. This means we should instead target the JSX element node so we have easy access to both properties.

[06:18 - 06:28] This is a simple case since the JSX element node is the closest shared ancestor and the parent of both nodes we need to transform. Keep in mind that this might not always be the case.

[06:29 - 06:39] The closest shared ancestor could be anywhere higher in the tree and often is not the parent node. For example, the closest shared ancestor node could be the grandparent of some nodes in the parent of others.

[06:40 - 06:48] Keep in mind that this is just a general guideline for determining which node to target in the tree. The node to target depends on the exact case.

[06:49 - 07:01] For complex transformations, it may even require targeting multiple nodes. With this in mind, let's switch back to our editor and instead target JSX element nodes.

[07:02 - 07:22] You'll notice we're now getting a type error since the name property doesn't exist on the JSX element node. We can now destructure the opening element and closing element from this JSX element node.

[07:23 - 07:33] Now that we've made this change, we can also mutate the closing element. We're not getting a type error because the closing element is possibly null or undefined.

[07:34 - 07:45] This is because some JSX elements can be self-closing and don't have a closing element. In our case, all of the button elements do have a closing element, but let's add a conditional to prevent this type error.

[07:46 - 08:04] We can now try running the script again. However, one thing to keep in mind is that each time this transform script is run, it will directly transform and rewrite code in the flashcards codebase.

[08:05 - 08:14] Assuming the transform is correct, this is the desired outcome. It's possible that while testing the transform, like we have been doing, we'll want to run it on the codebase several times.

[08:15 - 08:23] Before it can be run again, the previous changes will need to be undone. For now, the easiest way is to re-download the flash app from a previous lesson .

[08:24 - 08:35] Alternatively, you could manually undo all of the changes. A quicker approach would be to use git and run git resethard or an equivalent command to undo the changes from the transform.

[08:36 - 09:06] This will remove all your uncommitted changes, so make sure any previous work is committed. Now that we have a fresh flash codebase, in an easy way to undo transforms, we can switch back to the transform directory and now run the transform script again.

[09:07 - 09:18] Using the login component again in the flashcards codebase, we can see that this button element was properly transformed, including both the opening and closing element tags. This fixes our first issue.

[09:19 - 09:30] The next issue is transforming the props. When we implemented the button component in the previous module, defaults were added for some props, and some props were replaced with others.

[09:31 - 09:41] So far, we've renamed button elements to button components, but I haven't accounted for any of these other differences in the props. The next step is to update the props and make them valid with the button component.

[09:42 - 09:51] For props with values that match the new defaults, they can be omitted since they are redundant. For props that were replaced, they need to be converted to the replacement prop and the old prop removed.

[09:52 - 09:59] Lastly, there are some props that were unchanged so we can pass those through. Let's update the script to handle all of these prop changes.

[10:00 - 10:13] The specific changes are for the type prop and now defaults to button, so it only needs to be kept for places that set a different value, which would be submit. For the variant prop, we can set that to the button variant defined in the class name string.

[10:14 - 10:21] For example, button - secondary would be secondary. If it's primary, it doesn't need to be set since that's the default value.

[10:22 - 10:32] For the block prop, it only needs to be set if class name contained button block. For the Encla Candler, it's unchanged so the same handler can be used.

[10:33 - 10:42] Let's start by creating an array to track the props to be used for the transformed button component. Let's then give this array a type of JSX attribute nodes.

[10:43 - 11:03] To do this, let's import the types from the Babel types package. This Babel types package provides helpers to build different node types within the tree, but it also provides type definitions for those nodes, so it can be used for both.

[11:04 - 11:17] In this case, we're using the JSX attribute type, which will be an array of these types of nodes. Let's now loop through all of the existing props and transform those into the new props.

[11:18 - 11:45] Again, we're only expecting attributes that are JSX attribute nodes with a name that is JSX Identifier node, so let's add that conditional. Now, depending on the exact prop, we want to handle it slightly differently.

[11:46 - 12:10] Let's start by adding a switch statement with a case for each prop. The type prop only needs to be kept if its value isn't button, so let's add a conditional for this case, and if it's not button, add it to the array of new props.

[12:11 - 12:21] The class name prop is going to be the most complex transformation. To start, we expect it to be a string literal, so we can add a conditional for that.

[12:22 - 12:33] The class name value is a single string with all of the class names separated by spaces. Let's split this on spaces so that we have an array of all the class names.

[12:34 - 12:48] For the buttons variant, we know that this is determined by the class name that starts with button dash dash. However, there's also the button dash dash block class name that we want to ignore, so let's also add a conditional for this when finding the class name.

[12:49 - 13:06] And finally, we want to remove button dash dash, and that will leave us with the variant prop value. Similar to the type prop, we only need to set the variant if it's not equal to primary, so let's add that conditional.

[13:07 - 13:17] However, this case is a little different, because there's not an existing node to reuse. We need to create a new JSX attribute node for the variant prop and value.

[13:18 - 13:23] For this, we can use the Babel types package. It exports builders for each type of node in the tree.

[13:24 - 13:53] This helps ensure that all of the nodes are constructed correctly. [ Silence ] Notice how the builders start with lowercase letters, while the types start with capitals.

[13:54 - 14:04] If you're unsure what the parameters to a builder are, look at the type signature. Here we can see the JSX attribute builder expects a name, as well as an optional value.

[14:05 - 14:20] String literal expects a string, however, variant can be undefined, so let's update our conditional to also make sure variant exists before constructing this node. Now that we've handled our variant, the last thing to do is handle the button block class name.

[14:21 - 14:39] This is much easier since we only need to look for that exact class name, and if it exists, add a new JSX attribute node. However, this time is also slightly different, since it's a boolean, we don't need to explicitly set a value since it can act as a prop flag.

[14:40 - 14:54] And finally, for the onclick prop, we can pass through the existing node. We now have an array of the new props, so let's assign those to the opening elements attributes.

[14:55 - 15:10] Let's switch back to our terminal, and first undo any changes in the flashcards codebase. Then we can switch back to the transform directory and run the transform again.

[15:11 - 15:25] Now let's switch back to the login component in the flashcards codebase. This script now converted all button elements to a button component, and transformed the props to be compatible with the button component, and avoided any unnecessary props, which match the defaults.

[15:26 - 15:34] The script is almost able to fully transform the code. The final missing piece is the import for the button component, which will need to be added to the AST when necessary.

[15:35 - 15:48] Before we update the script, pause the video and take a moment to paste some code with several imports into AST Explorer. We can see that all the import nodes are at the beginning of the program body array.

[15:49 - 15:57] The simplest approach should be to add it to the front of this array. A more complex heuristic, such as alphabetically by package name, could also be used.

[15:58 - 16:10] For now, let's just add it to the start of the file to make things simple. A single file could contain multiple button elements and therefore multiple button components, so we only want to add one import per file.

[16:11 - 16:20] So let's track the button usage on a per file basis. Then if we found a button, we can switch this boolean to true.

[16:21 - 16:32] And finally, if the file did contain a button, let's add an import. The first step is to construct the import path.

[16:33 - 16:44] This project is using relative imports, so let's continue to use that. We can use node's path.relative helper to determine this path dynamically.

[16:45 - 16:59] The path.relative helper expects two arguments, the from and to path. In this case, we're using another helper to get the directory name of the file that we're currently transforming.

[17:00 - 17:31] Now that we have the import path, we can construct a new import node. We can see that the import declaration builder expects an array of specifier or imports, as well as the source, in our case, the relative path to the button component.

[17:32 - 17:47] Now that we have a new import node, let's add it to the top of the program body . After resetting the previous transforms changes in the flashcards codebase, run the script again.

[17:48 - 18:02] Looking at the logging component again in the flashcards codebase, we can see that the button was still properly transformed, but this time, we also now have an import for the button component with the proper relative path. This script now solves all of our requirements.

[18:03 - 18:12] We transformed button elements to button components. We transformed button props to button compatible props, and finally edited import for the new button component.

[18:13 - 18:19] One other thing you might have noticed is that the formatting changed. The formatting changed because the original code is converted into an AST.

[18:20 - 18:31] The AST won't retain the formatting, so when converted back to a code string using the generate function, the output can have different formatting. Much of the flash app codebase is already formatted using prettier.

[18:32 - 18:39] It's possible that this transform script could be run first, then run prettier after that. However, it would be nice to run it all in one go.

[18:40 - 18:56] Fortunately, prettier provides an API to do so, which can be used to format the source code string before writing it back to the file. Let's add prettier and its types as dependencies for our script.

[18:57 - 19:15] Now let's switch back to the transform script and import prettier. Then right before you write the new source code string, let's format it.

[19:16 - 19:28] With a few lines of code, we've now updated our transform to format the output code consistently with the rest of the project. Interestingly, prettier itself is using Babel parser for formatting the code.

[19:29 - 19:42] Another example of everyday tooling that relies on ASTs. After resetting any changes in the flashcards codebase, let's run the script again to confirm it still runs properly.

[19:43 - 19:57] After running the script, you may notice slightly fewer formatting changes as a result of running the transform. This script is now able to complete all of the steps needed to transform from button elements to button components, including the transform props.

[19:58 - 20:09] To quickly review, the transform script was adapted directly from the audit script and reused much of its logic since the AST and nodes which we wanted to target were similar. However, we ran into a few problems.

[20:10 - 20:15] First, only the opening element was renamed. Second, the props needed to also be transformed.

[20:16 - 20:21] Third, the button component wasn't imported. And fourth, some unexpected formatting changes occurred.

[20:22 - 20:31] We were then able to tackle each of these pieces individually. First, instead of targeting only the JSX opening element, we instead targeted the parent, the JSX element.

[20:32 - 20:42] This allowed easily modifying both the opening and closing elements. Second, the script looped over the existing props and either passed those pops through or transformed them into new props.

[20:43 - 20:55] Third, if there was at least one button used in the file, a new import was added at the very top of the file. And finally, prettier was run on the source code string output from Babel Generator to prevent unexpected formatting changes.

[20:56 - 21:06] At this point, you might start thinking that this took a while or required a lot of boiler plate or that it might not scale to thousands of files. Fortunately, JS code shift is a specialized tool for transforming code like this.

[21:07 - 21:10] In the next lesson, we're going to recreate this script with JS code shift.