Practical Refactoring with Syntax Trees

Speaker 1 [00:00]

So, good afternoon everybody. Please find your seat. We are about to start in a couple of seconds. It's my great pleasure to introduce our next speaker, Laurent Deray. He will talk about refactoring Python code by using the super power of abstract syntax trees. Please give him a warm hand of applause.

Speaker 2 [00:42]

Thank you. Hello, everyone. I'm Laurent, and we're going to talk about syntax trees and refactoring. So we're going to talk about what's an abstract syntax tree. Because before we can go into automated refactoring, we need to have a way to represent our code as data that we can manipulate. We're going to look at defining refactoring scripts, what are we aiming for? And then we'll walk through an example using libcst, it's a library from Instagram, to refactor code with syntax trees. And I'll share some tactical advice on how you can maybe get started writing such scripts. So before we go into syntax trees, let's look at what we want to achieve. A refactoring script is a script that takes your files on disk, transforms them, and then writes them back to disk, hopefully improved. So this is an example of the end game. This is what we are targeting, basically, this bash script. It's already a refactoring script. So when regex works, I think it's a great tool to use. But sometimes they fall short, and we can look for more sophisticated tooling. So let's start with the tree part of abstract syntax tree. If we look at this expression like this, it can be represented as a tree. So there's an operator, a multiplication. There's a left-hand side, a right-hand side. And the left-hand side is itself an operation. And it can convert to this tree. So all the nice visualizations were part of the tutorial this morning, the tutorial session on ASTs. If you missed it, you're stuck here, I think. So this tree, we have the root node. It's a multiplication operator. And two children, a subtree on the left, and just a constant on the right. And if you had to turn yourself into a computer, you might evaluate the expression by starting at the bottom of the tree and doing 1 plus 2, and then climbing back up the tree to evaluate the expression. So if we look at what Python does in more detail with such an expression, we can put it into a file and ask the AST standard module to show us what it looks like. So there's a module. It's always a top-level element in the Python module. It has a body. There's only one item in the body. It's an expression. And inside the expression, it's the core of our code here. It's what we've seen with a binary operator, a left-hand side, a right-hand side, and a multiplication operator. So I'm not claiming that this is easy to read or very user-friendly. But this is what we are going to use to transform code. So every node is a Python object. And it has children or attributes on the object. And it's all Python, so we can all access and modify it. So what is the AST? The AST is a tree of Python objects that represent the source code. Node types represent language constructs, so expression, assignment, import, function definition. Everything, every syntax in Python has a matching node type. It's a useful representation of the code, especially, I mean, it's used by the Python interpreter to actually execute Python code. It's not super useful to read the code. So if you look at the AAC, it's usually less human-friendly than just text. But it's a representation of the code that we can modify. So if we look at what this looks like with the standard library, we can import AST, take some source code, pass it, this gives us a tree, a syntax tree. It's just a Python object that we can dump to a string like we did earlier. And we can also unparse it. So we can convert the tree back to source code. Now, unparse we'll see later, but it's mostly a debugging feature in the ASC module. But it gives us an idea for what refactoring script could look like. So we could take the code, read the code, pass it into a syntax tree, transform it into a new syntax tree, and then write that new syntax tree back to disk. turn it back into code, and write it back to disk. So in code, this is really the whole pipeline that will be like the refactoring script that we are going to look at. Read the source code, pass it, transform the tree, and write it back to disk. So this is a whole pipeline, but the real logic lives in this line, so the tree transformation. And this is what we are going to focus on next. So transforming the abstract syntax tree. The module gives us utilities, nodeVisitor, nodeTransformer. You're maybe familiar if you've been to this morning session. They help you traverse the tree and mutate some nodes, only the nodes that you'll either look at some nodes or mutate some nodes that you are interested in. So depth first traversal. And as a user, what you do is you define visit method for the node types that we are interested in. Like if we want to look at all the name nodes, we define visit name. All the function definition nodes, we define this method. And it will be called by the node transformer. So let's look at an example from the docs. They show an example to transform the code on the left into the code on the right. If you had to do this yourself, I mean, if this was your goal somehow, I think the first thing would be to define what we want to achieve. So we know what we want to achieve in terms of code. What does it mean in terms of abstract syntax tree? Thanks. So if we look at the trees, again, not super user-friendly, but we can identify what we want to change in this, and what needs to change is the name nodes that are the target of the assignment or on the left of binary operations, all the name nodes, they need to become subscript nodes. So this subscript is how Python describes the indexing operation here, and we can look at what we want the tree to look like in the after state, and this is how we are going to define our transformer. So we inherit from node transformer, we define only one method, visit name nodes, and we return the new node that we want to have instead. So in these few lines of code, it can already be a little mind-bending if you try to think about it, because the slice we're including here is a constant with the value that is the name of the variable, so it can take some time to process, but what's important is the idea for now. So let's say we wanted to use the Python AST for refactoring, there's a bit of an issue because these two snippets of code have the same AST. So if we pass the one on the left and we dump it back to code, we're going to get the one on the right. So we're going to lose the important comment, we're going to lose the parentheses because these parentheses in this case are just decorative. So the computer does not care, and it does not remember that they were there. So that makes it, I mean, not ideal fit for refactoring script. I will say that some tools still manage to use that, like PyUpgrade and DjangoUpgrade. They use the AST mostly to identify the patterns that they want to change, and then they use some other techniques to actually do the rewriting. So it does not preserve formatting. It does not preserve comments. So basically not ideal. Good thing is we have concrete syntax trees. So this is not exactly the same example as on the previous page, but close enough. So concrete syntax trees, essentially the same thing. It looks pretty similar, right? You have an expression, a simple statement line, an expression, and binary operation again with left, right, operator. So same ID, only there's some extra information. There are parentheses, and there's white space with comments. So this is going to help us preserve the comments when we are refactoring. So concrete syntax trees are cousins of ASTs. Sometimes they're called parse trees, sometimes AST and CST. is a bit blurry, so the terms are used interchangeably. And so there are different trees because there are different libraries, so the note types might not exactly match, but the mental model of working with syntax tree transfers. And so it helps preserve white space, parenthesis, comments. It allows round tripping, so that's like writing back exactly what we got in after parsing and un-parsing. So it makes it a great fit, and we do not really have to care about white space or comments when writing the scripts, it just has to be here for the pipeline to preserve them, but we don't have to manipulate them, really. So let's look at some, get a feel of what the tree looks like, similar to the abstraction text tree, but seeing more of it helps. So this is going to be an import node and a list of names. So you can already see maybe it's also a way to represent multiple imports. Import alias, and ultimately, a name node with a value numpy. So now if we add an alias, we see that the import alias gains an attribute that's as name. And oh, sorry. And yeah, so all the information is in the tree. So we also have assignments. Here we are just assigning to a string to A. And if we add a function call in there, it just adds a node in the tree. So we're going to look at an example now of how we can use this concrete syntax tree and transform it to just rewrite code. So the example we're going to take is PyTest fixtures rewriting. So PyTest fixtures, PyTest is a testing framework. And PyTest fixtures are a way to set up data. So they're a little bit magic in how they work, in that in this snippet, test login is the test. And when it has parameters, PyTest knows to look up the name of the parameter as a test fixture. So it sees that test login has test user as parameter. It figures it has to be a fixture. So it looks up the fixture function, and it runs it. So what's wrong with this code? If you're familiar with PyTest, you might be a little upset, because this is not the correct naming convention for a fixture. So the fixtures start with test underscore, which is the naming convention for tests. So it can be difficult to look at, I understand. And also, it can cause crashes with some setup, some PyTest versions. PyTest is pretty good about making the best out of it. But under some circumstances, it can cause crashes. So now, say you have hundreds of fixtures like this one. How do we change this? If there's just one, you go into your IDE, and you rename, and it's easy. If you have hundreds of fixtures like this one, it's a bit trickier. It becomes a very repetitive task. So let's try to write a script to do this. And we're going to start with what we want. So we have the before code with the terribly named fixture. And we want to just rename it to user fixture so that there's no ambiguity for what is a test and what is a fixture. So collecting these snippets when starting to write a code mode really helps understanding edge cases and if what you want to achieve is well-defined, and they can also become test cases, so it helps you write the code mode eventually. So what does it take to do this transform? We're going to need to identify the test fixtures and then rename them everywhere they are used. So yeah, this is what changes from A to B. We rename the fixture, the test user, into user fixture. So we're going to use libcst for that. So there's a CST transformer. It's similar to a node transformer. We define methods like a visit underscore node type and leave underscore node type. The visit method is essentially going top down and reading, traversing the tree. And the leave method is where you can get an opportunity to return a new node, so a transformed node. We are going to look at this transformer. We know that we want to collect the PyTest fixtures, so we are going to need to visit function definitions and identify PyTest fixtures, and then we are going to need to transform some variable names, so using the name node. We can transform, yeah, we can, when visiting a function definition node, we can look at whether it's a PyTest fixture node, and we should rename it, and then we collect the name to rename into some dictionary that we just store on the class. So it's a bit of wishful programming here, because we don't have the isPyTest fixture function yet. But once we have that, so we collected all the renaming that we needed to do, we can define a leave name method and update all the variables that match fixtures that we want to rename. So this is relatively straightforward code when it's not on the slide. we return a new node that uses a renamed variable. So the big part that remains is matching fixtures. So how do we identify that a function definition node is defining a PyTest fixture? To do that, we're going to look at some more trees. So if we look at the regular function definition, it looks like this. Function definition has a name, that's the name of the function. And if we add a decorator, the decorator's attribute gains an item, and we can see what it looks like. So maybe to match fixtures, all we have to do is try to match this structure. So the pattern matching in Python helps with that. And we can look at all the decorators, and if one of them looks like PyTest fixture, then this is a PyTest fixture. And if none of them are such decorators, then we don't have to rename this function. And that's basically all there is. And with this, you can rename as many tests underscore fixtures as there are in your code base. And that's all it takes, basically. So now I am going to share some practical advice and then we are going to come back to this example that has some serious limitations. So this is a very serious slide about running code modes and you should be strict and rigorous when running them for your own good. So you want to start with a clean Git working tree. Make sure the Git status is clean. Run the code mode. So if something goes wrong at this stage, you can collect a new test case and just discard all the changes. I like to run formatters and linters separately as a separate step. And then you commit just the automated changes. So it's a bit of boring advice, but I'd feel bad if I didn't mention this, because when you are going to rebase your pull request on top of new changes and you have conflicts, it's easier if you can just drop all the automated changes, avoid conflicts, and then run the script again. If you mix automated and manual changes, you're a bit in a pickle. So now writing code modes, what do we care about? Maintainability, not really, because the point of the code mode script is just to run once, automate work that you would maybe have done manually otherwise. So you don't even have to check it into Git. Maybe you just want to pretend to your coworkers that you're very busy and just make the pull request appear as if you've done them manually. So maintainability doesn't really have to be a concern in all cases. Sometimes it's okay to not go all the way as well. Like some edge cases can be very difficult to deal with. but if you automate 90% of a change, maybe it's good enough. Leaving formatting to tooling really helps making it easier because you don't have to manipulate white space in your code mode. You don't have to make the output nicely formatted. You can just run black or rough afterwards. And it also works with duplicate imports. If you have something to duplicate afterwards, then your refactoring script can just add as many imports as it thinks is necessary. So, yeah, in this test-driven development, I don't know if it's a positive thing or not, depends on people, but it's a really great fit for refactoring scripts. So yeah, you might notice that this combination of, like, not completely successful, not maintainable Test-driven development is a good fit for AI agents sometimes, especially the 90% thing. So if we go back to our example, there are some serious limitations. So variables that are the same name as fixtures but are local variables would be rewritten. So that's not desirable. So I really followed my own advice about claiming success. Sometimes it's good enough, like on the codebase that I was working on, this case didn't happen, so I didn't have to go and fix it. There's another issue when the fixture is defined after the place where it's used in the code, because the tree more or less keeps the order of the source code, if the function definition is defined after it is used as a fixture, it will not be detected by the code mode. And finally, fixtures across files are not supported here. So also matching is not the most robust, but we will see what we can do about this. So if we look quickly at one example, maybe defined after use, in this case, the fixture definition is after the test function, and this is a failing test case for our code mode. But if your tests have been properly copy pasted thousands of times, then you might not have this issue. So how do we go further with libcst and maybe address some of these shortcomings? There are ways of addressing the state across files issue using multiple passes, so first a pass on all values to collect data, and then a second pass to transform it. So this uses the metadata APIs, and in particular, there's a scope provider utility that helps manage the scope of variables. So this would help with our shadowing issue so that we can identify that a variable that has the same name as one that we want to transform is only a local variable and not an actual test fixture. And our utilities to pattern match nodes in more reliable ways. So how would you get started with automated refactoring? The main point I want to make is that it doesn't have to be super mysterious, and it can be accessible, and already understanding that it exists can help a lot. You can use existing tools. You can use Django upgrade, PyUpgrade, or you can use, if you're a Next.js user, you can use Next.js official code modes to upgrade to the Canary version, which is, I think, recommended practice. And then you can start asking yourself when you have boring repetitive work, would this be doable as a refactoring script? Can I automate this? I will say that often for me it's a lot more fun writing a refactoring script than it is doing three hours of repetitive work. We can look at, we have time I think for a couple of ideas for refactoring scripts. So this is things that exist. So you could convert unit test classes to pytest. If you've been told to change testing framework, there's a way to maybe automatically rewrite this to this version, so there's an open source project that does not use libcst and is actually a good candidate probably for a port because it doesn't support recent Python versions. This one we're going to skip maybe, but here's another ID, clean up feature flags automatically, So maybe you have feature flags everywhere in your code base. This is JavaScript, I apologize. Maybe you have feature flags everywhere in your code base, and you can have a script to say, okay, remove the new release feature flag, because that was like four years ago, and we are done now. So you can automate this. I put the blog post that mentions this particular example, but you can automate this as well in JavaScript. So now for the last bit of the talk, I'm going to address the elephant, AI, of course. Why would you write a libcst script when you can just ask AI to change this and do that? So this is a bit maybe of a biased view, you'll take what you want from it. So one is deterministic, one is based on your face of agents, and I don't mean this necessarily in a bad way, but it means that every time you rerun such a transform, you have to check the output again if you do not really trust it. So to make sure like Claude did not get creative on some file. So yeah, with deterministic script, you can rebase easily, or you have to like run everything again with Claude. If you really insist, Claude can write the the libcst-transformer, it's pretty good at doing, I mean, I would say one-shotting the very basic things. Then building on top of that is probably a little harder, but as a first try, if you have all the test cases, AI does pretty well when it has test cases to iterate on. So the clear and obvious conclusion is that AI is still really good in many cases. So if you have a smaller code base, or you do not feel like you are going to have to rebase this, it is just some one-off change, it is pretty good at making code base-wide changes. And code modes feel better for some changes. So you have to maybe make your own decisions on this. Yeah, and the larger your code base, and the more you are going to need to reuse maybe the script, the more you are pushed towards code modes and refactoring scripts. So, yeah, I want to end by saying, like, the syntax tree, the abstract syntax trees, concrete syntax trees, they sound a little bit mysterious sometimes when you hear about them. Sometimes it's also because they're used interchangeably and in weird contexts. But when you look at it, it's actually quite approachable, and it gives you a lot of power, And also, it gives you an idea of all the, I mean, how, it helps you understand how the dev tools that you use every day work under the hood. So yeah, some resources on the slide about me, but we don't have time for marketing. Thank you. And come talk to me if you're...

Speaker 1 [29:29]

Thank you very much for this presentation. We have a few questions. The first question is, is AST-CST refactoring actually used by IntelliJ PyCharm when user does refactoring?

Speaker 2 [29:46]

So, I don't know about IntelliJ and PyCharm in particular. I suspect that they developed their own tools because LeapCSC is not that old. But you can achieve similar things that they do. LeapCSC is from Instagram, and they used it to power linting and refactoring on their codebase.

Speaker 1 [30:10]

Next question is, can you recommend any tools for quickly visualizing the CEST IDE plugins, web tools or something?

Speaker 2 [30:19]

So for the CST, I don't really know. I mean, there are utility functions that just dump the tree in the library, so that's good. For the AST, I think there's one in the resources. Yeah, there's ASTExplorer.dev, which is really a great way to explore ASTs in many languages. So for Python, it's using Pyodide To construct this Python in WebAssembly, you put Python code and it puts out the actual Python AST.

Speaker 1 [30:58]

Okay, thank you. We would have maybe time for two questions from the audience.

Speaker 2 [31:17]

Hey, thank you for the talk. How would you say did you go about learning about ASTs? Was it a practical experience or more theoretical? So I think I first learned basically out of curiosity. I started with the PyTest assert statement. So if you have a plain assert statement, Python gives you a traceback that just says assertion error, something broke. And if you want to know why A was not equal to B, you have to run again with print statements. So I wanted to know how that worked. So that's how I looked into ASTs in the first place. So they actually transformed the AST. And they do not write to disk, but they want to make the code run as if it was a little bit different. So they transformed the tree before passing it back to Python for execution.

Speaker 3 [32:14]

You told AST is used by Python itself, or did I get it wrong?

Speaker 2 [32:23]

No, the AST, you got it right. The AST is used by Python to execute code.

Speaker 3 [32:31]

Just for pre-parsing, because you told it only knows the type at the end if it's defined.

Speaker 2 [32:31]

Just for...

Speaker 3 [32:40]

Because if it's used as a fixture, as a parameter at the top, it will only later know this is a fixture. So how can Python use it if it's not going over many files? or is it just because you just put this piece of code? I mean, there must be something above it because otherwise it's only a parser.

Speaker 2 [33:07]

Yes. So the AST is still a representation of the code, but it does not have all the information that comes from executing the code. So Python uses it. I'm going to try answering and then you'll tell me if I got it correctly. So Python takes the source code, it tokenizes, parses it into abstract syntax tree and then compiles it into bytecode. And that is what is interpreted at the end. Only pre-parser.

Speaker 3 [33:33]

It's a pre-puzzle tree. Sorry? It's only a pre-puzzle tree.

Speaker 2 [33:37]

Yeah, I mean, you could say, like, usually in the naming, there's, like, parse tree, concrete syntax tree, and abstract syntax tree is the most abstract. But, yeah, it's a product of the compiling process, essentially. And the scope of variables is resolved at runtime.

Speaker 1 [33:59]

So we have