Automated Code Migrations: Our Journey from Unittest to Pytest

Photo by Patrick Tomasso on Unsplash

Code migrations can be a challenging task. Even when executed correctly, they often result in:

• An initial wave of easily migrated files
• An expanding documentation detailing the migration steps for the remaining codebase, which grows with each discovered edge case
• A decrease in enthusiasm and motivation as the complexity increases, leading to a slowed or halted migration process

Does this sound familiar? Well, we found a way out of that.

Enter Codemods

Codemods are tools that automate the process of making changes across a codebase, which would be tedious and error-prone if done manually. They are particularly useful when dealing with large codebases and repetitive, predictable changes.

Codemods work by parsing your source code into a Syntax Tree, applying transformations to it, and then generating a new source code from the updated tree. This process allows for precise, syntax-aware transformations of the code.

It’s just code

Just like any other code, codemods can be iteratively refined and tested to ensure they are applying the desired transformations correctly. Instead of an ever-growing list of migration instructions, which can slow down the process and increase the risk of manual errors, handling a new edge case with a codemod is as simple as adding a test and updating the code to make it pass.

Moreover, codemods promote better distributed ownership. By automating repetitive parts, they enable more team members to participate in the migration process. This not only distributes the workload but also fosters a sense of shared responsibility over the codebase.

Codemods in Action: Migrating from Unittest to Pytest

Let’s illustrate with a recent migration where we changed our test framework from unittest to pytest. We leveraged codemods to automate the migration using libCST, a Python library for parsing and manipulating Syntax Trees in Python.

Here’s an example of a codemod allowing to change self.assertEqual(a, b) unittest assertions to their assert a == b pytest counterpart:

import libcst as cst
import libcst.matchers as m


# We implement our own cst.CSTTransformer subclass
class UnittestToPytestTransformer(cst.CSTTransformer):
    # We only want to transform a call expression
    # this extension method will be called on each Call node of the tree
    def leave_Call(
        self, original_node: cst.Call, updated_node: cst.Call
    ) -> cst.BaseExpression:
        # Using matchers (think regex but for a syntax tree)
        # we target only the `self.assertEqual` calls
        if m.matches(
            original_node,
            m.Call(
                func=m.Attribute(value=m.Name("self"), 
                attr=m.Name("assertEqual"))),
        ):
            if len(original_node.args) < 2:
                return original_node  # Not enough arguments, return as is.

            # If there's a third argument (an assertion message)
            message = updated_node.args[2].value if len(updated_node.args) > 2 else None
            new_node = cst.Assert(
                cst.Comparison(
                    left=updated_node.args[0].value,
                    comparisons=[
                        cst.ComparisonTarget(
                            operator=cst.Equals(),
                            comparator=updated_node.args[1].value
                        )
                    ],
                ),
                message=message,
            )
            return new_node
        return original_node

To run this codemod on a Python file, you can use the libcst.tool command-line interface provided by libCST:

python -m libcst.tool codemod UnittestToPytestTransformer directory

This command applies the UnittestToPytestTransformer codemod to any Python file inside of the directory, replacing it with the transformed version.

Conclusion

Codemods are a powerful tool for large-scale codebase refactoring, making changes easier and more efficient. Their ability to be improved and tested, and their contribution to distributed ownership, make them an invaluable tool for any software engineering team.

In our case, we had been working on the migration from unittest to pytest for years, involving many engineers in the process. This tool allowed us to complete this migration within a quarter, thanks to a team of 2 dedicated engineers developing the codemod, and a handful of volunteers applying it to the parts of the codebase they were most familiar with.

This initiative not only improved our codebase but also demonstrated the effectiveness of codemods. In the future, we will continue to leverage them in our refactoring efforts to improve our codebase quality and team productivity.