Practical Refactoring with Syntax Trees
Hello everyone
I'm Laurent.
We're going to talk about Syntax Trees.
And refactoring.
Before we can go into automated refactoring, we need to have a way to represent our code as data.
Practical applications: why you should care.
Example:
#!/usr/bin/env bash
sed -i 's/old_function/new_function/g' *.py
Let's define refactoring scripts
Before "example":
As a spoiler, here's an example of the "endgame", what we're targeting...
When the basic tools work: use the basic tools.
Use the IDE refactoring, use regexes.
Now regexes quickly show limits. Multi-line patterns are a pain, they don't "know" about the language,
etc.
I'm sure I don't need to convince you.
To do better than regex, we are going to use another way to represent the code. Instead of seeing the
code as text, we'll see it as a datastructure, a Syntax Tree.
Code as Data
(1 + 2) * 3
Let's start with the "Tree" part of Abstract Syntax Tree.
An expression like this one can be represented as a tree.
There's an operator. A left-hand side, right hand side.
The left-hand side is itself an operation.
It converts to this tree. The nicer visualizations were all in the tutorial session this morning.
If you missed it, you're stuck with my design skills.
DESCRIBE THE TREE.
- Root node: multiplication operator
- left hand side and right hand side as children.
- subtree as child for left hand side.
If you had to turn yourself into a computer, you might evaluate it starting from the bottom: you'd do
1 + 2 first - essentially evaluating the left child subtree - then you'd use that to evaluate the root
node.
$ echo "(1 + 2) * 3" > simple.py
$ python -m ast simple.py
Module(
body=[
Expr(
value=BinOp(
left=BinOp(
left=Constant(value=1),
op=Add(),
right=Constant(value=2)),
op=Mult(),
right=Constant(value=3)))],
type_ignores=[])
Let's take a look at the Python AST for this expression.
Now there's more stuff here.
There's a module. Basically always the top level element.
The module has a body, which has a list with just one expression.
There we find a tree similar to the one we had.
Notice we got that tree from running the `ast` module from the standard library as a script with the
`python -m` syntax.
You can run that in your terminal on your Python files. It's not particularly readable - we're better at
reading source code.
Every node is a Python object, and the children here are attributes on the object.
There are nodes for FunctionDefinition, Assignment, For loops...
You don't need to know them beforehand.
What is the Abstract Syntax Tree?
a tree of Python objects that represent the source
node types represent language constructs: expression, assignment, import, etc.
used by the Python interpreter to run code
a representation of the code that we can modify
So what is the AST?
Independent of formatting: no whitespace, no comments. We lost some information, can't reconstruct the
original source code exactly.
A useful representation of the code.
Useful for some purposes. NOT reading the code.
>>> import ast
I said the AST was a useful representation. How do we use it?
It's part of the standard library
>>> import ast
>>> source = "(1 + 2) * 3"
Start with some source code
>>> import ast
>>> source = "(1 + 2) * 3"
>>> node = ast.parse(source)
ast.parse: turn python code into an AST.
>>> import ast
>>> source = "(1 + 2) * 3"
>>> node = ast.parse(source)
>>> node
<ast.Module object at 0x7bcbef830370>
We have a module object. It's a Python object with attributes on it.
>>> import ast
>>> source = "(1 + 2) * 3"
>>> node = ast.parse(source)
>>> node
<ast.Module object at 0x7bcbef830370>
>>> ast.dump(node)
'Module(body=[Expr(value=BinOp(left=BinOp(left=Constant(value=1), op=Add(), right=Constant(value=2)), op=Mult(), right=Constant(value=3)))], type_ignores=[])'
We can dump the tree to a string representation
This is what the ast module does when invoked on a file.
>>> import ast
>>> source = "(1 + 2) * 3"
>>> node = ast.parse(source)
>>> node
<ast.Module object at 0x7bcbef830370>
>>> ast.dump(node)
'Module(body=[Expr(value=BinOp(left=BinOp(left=Constant(value=1), op=Add(), right=Constant(value=2)), op=Mult(), right=Constant(value=3)))], type_ignores=[])'
>>> ast.unparse(node)
'(1 + 2) * 3'
We can also convert the tree back to code.
Now Python does not need that because it can execute the AST.
If we modify the AST, we can give it to Python and it can execute it.
But this is very convenient to debug (new in Python 3.9).
Anatomy of a refactoring script
# read, parse, transform, write
source_code = read(file_path)
Now the topic of the talk is refactoring.
So let's see where syntax trees fit in.
Anatomy of a refactoring script
# read, parse, transform, write
source_code = read(file_path)
tree = parse(source_code)
Anatomy of a refactoring script
# read, parse, transform, write
source_code = read(file_path)
tree = parse(source_code)
transformed_tree = transform_tree(tree)
Anatomy of a refactoring script
# read, parse, transform, write
source_code = read(file_path)
tree = parse(source_code)
transformed_tree = transform_tree(tree)
write(transformed_tree.unparse(), file_path)
Anatomy of a refactoring script
# read, parse, transform, write
source_code = read(file_path)
tree = parse(source_code)
transformed_tree = transform_tree(tree)
write(transformed_tree.unparse(), file_path)
Transforming the AST
Utilities: ast.NodeVisitor, ast.NodeTransformer
Depth-first traversal
Define visit_[NodeType] methods
ast.NodeTransformer
after.py
data['b'] = data['a'] + 1
Assign(
targets=[
AssignTarget(
target=Name(value='b')
)
],
value=BinaryOperation(
left=Name(value='a'),
operator=Add(),
right=Integer(value='1')
)
)
Assign(
targets=[
AssignTarget(
target=Subscript(
value=Name(value='data'),
slice=[
SubscriptElement(
slice=Index(
value=SimpleString(
value="'b'"
),
)
)
],
)
)
],
value=BinaryOperation(
left=Subscript(
value=Name(value='data'),
slice=[
SubscriptElement(
slice=Index(
value=SimpleString(
value="'a'"
),
)
)
]
),
operator=Add(),
right=Integer(value='1')
)
)
ast.NodeTransformer
after.py
data['b'] = data['a'] + 1
class RewriteName(NodeTransformer):
def visit_Name(self, node):
return Subscript(
value=Name(id='data', ctx=Load()),
slice=Constant(value=node.id),
ctx=node.ctx
)
Now say we want to modify the AST.
Let's look at this example from the docs.
ast module gives us The NodeTransformer class
We can use it to transform an AST.
This is when we have a piece of code. And we want to modify it, as if rewriting it on-the-fly.
NodeTransformer It follows the "visitor" programming pattern, meaning we define the methods as
`visit_NodeType` and they
will be called on relevant nodes as the tree is visited.
The way you'd do this is you'd print both trees: before and after.
Then you know how to change
Subscript node means bracket access. The slice is the index or the key we use inside the brackets.
Replaces all variables like 'foo' with 'data["foo"]'.
The part we're interested in is the NodeTransformer. This is a pattern that can be used to modify the
AST.
Python AST for refactoring
exhibit A
1 + (2 * 3) # IP-protected formula
Same AST 😱
Same AST: if we parse the left, unparse it: we get what's on the right.
Lost comment, lost parentheses.
Because these parentheses are decorative. Such parentheses (and comments) are for humans.
The AST is a format for the machine :).
Python AST for refactoring: oh no
does not preserve formatting
does not preserve comments
Good if your codebase is not formatted and comments are overrated anyway.
Some tools still use the AST for refactoring, like pyupgrade and django-upgrade.
Two very successful tools, they use the AST to identify the patterns they want to change.
The rewriting is done differently, at a lower level.
MAYBE:
I'll mention it briefly: AST Transforms are still useful outside of refactoring scripts.
When you don't want to change the code on disk, but you want to make it behave as if it were different
code.
You can parse the code into a tree, transform the tree, and pass that tree back to Python for execution.
There are a few niche applications, pytest assert rewriting.
Concrete Syntax Trees
(1 + 2) * 3 # comment
SimpleStatementLine(
body=[
Expr(
value=BinaryOperation(
left=BinaryOperation(
left=Integer(value='1'),
operator=Add(),
right=Integer(value='2'),
lpar=[LeftParen()],
rpar=[RightParen()],
),
operator=Multiply(),
right=Integer(value='3'),
),
),
],
trailing_whitespace=TrailingWhitespace(
whitespace=SimpleWhitespace(value=' '),
comment=Comment(value='# comment'),
),
)
This is not even the same example as on previous slide.
But close enough!
Concrete Syntax Trees
Cousins of AST. Sometimes called Parse Trees.
Preserve whitespace, parentheses, comments...
Allow 'round-tripping'
Great fit for refactoring scripts
We do not have to care about whitespace
They are Cousins of AST.
Sometimes they are called Parse Trees.
Sometimes... they are also called ASTs. The frontier is a little blurry, but basically...
Different trees: the node types and structure won't be exactly the same.
But the mental model transfers.
We don't have to care about whitespace, and in fact, we will not.
We also won't really have to care about comments in our refactoring scripts.
We do care that they are preserved, but that is done essentially without us having to do anything
If you put all the formatting information aside, it looks similar to the AST.
Let's look at some examples now to get a feel of the tree structure and see different types of nodes.
Concrete Syntax Trees
import numpy
Import(
names=[
ImportAlias(
name=Name(
value='numpy'
)
)
]
)
View is slightly simplified, default values are not displayed, etc.
But that's the gist of it.
Concrete Syntax Trees
import numpy as np
Import(
names=[
ImportAlias(
name=Name(
value='numpy'
),
asname=AsName(
name=Name(
value='np'
)
)
)
]
)
Concrete Syntax Trees
a = 'Hello'
Assign(
targets=[
AssignTarget(
target=Name(value='a')
)
],
value=SimpleString(value="'Hello'")
)
Concrete Syntax Trees
a = f('Hello')
Assign(
targets=[
AssignTarget(
target=Name(value='a')
)
],
value=Call(
func=Name(value='f'),
args=[
Arg(
value=SimpleString(value="'Hello'")
)
]
)
)
Example: rename pytest fixtures
@pytest.fixture
def test_user():
return {"name": "test"}
def test_login(test_user):
...
What's wrong with this?
Naming convention
This is wrong
Why would you
Hurts my feelings
Can cause crashes with some pytest versions
Hundreds of fixtures like this one.
Right after showing the snippet:
EXPLAIN PYTEST FIXTURES. MAGIC PEOPLE MIGHT NOT KNOW.
Mention they can be in fixture arguments too (fixtures reference fixtures).
So if you're familiar with pytest you might be fuming at the sight of this snippet, because...
Naming convention!
It can cause crashes!
Yeeeeees
Now we have an excuse to do something about it!
If there are hundreds of fixtures like this, the IDE won't cut it.
It helps you rename one fixture: you rename it, it renames all the usages of the fixture.
Assuming you're using the correct IDE, but let's not get into it.
What do we want?
before.py
import pytest
@pytest.fixture
def test_user():
return {"name": "test"}
def test_login(test_user):
assert test_user["name"] == "test"
after.py
import pytest
@pytest.fixture
def user_fixture():
return {"name": "test"}
def test_login(user_fixture):
assert user_fixture["name"] == "test"
Change this, on hundreds of fixtures.
I like to start with before/after snippets.
What change do we want to make?
Writing the snippets (collecting from your codebase) helps understand edge cases/if what we want to
achieve is well-defined.
These should become test cases.
It also helps to think about what needs to be done: what does it take to do this transform?
First identify fixtures
Then rename them... everywhere where they are used.
LibCST API
CSTTransformer ~= ast.NodeTransformerDefine methods:
visit_[NodeType]leave_[NodeType]
We're going to use LibCST to transform our source syntax tree into the target syntax tree.
So, to be clear: the CST and the AST are not the same trees.
We've said that the CST has formatting information: commas, parentheses, whitespace...
But also the nodes might have different names and structure.
At the end of the day they're close enough though, the mental model persists.
not mentioned because you don't need to know the node types beforehand to write a transformer.
I don't know them
Just look them up as you need them.
import libcst as cst
class Transformer(cst.CSTTransformer):
def visit_FunctionDef(self, node):
...
def leave_Name(self, original_node, updated_node):
...
class Transformer(cst.CSTTransformer):
def __init__(self):
self.renames: dict[str, str] = {}
def visit_FunctionDef(self, node):
"""Collect fixtures that need to be renamed."""
if is_pytest_fixture(node) and should_rename(node):
old_name = node.name.value
self.renames[old_name] = generate_new_name(old_name)
return True # Continue visiting children
A bit of wishful programming. is_pytest_fixture, let's pretend we have that.
class Transformer(cst.CSTTransformer):
def leave_Name(self, original_node, updated_node):
"""Update variables that match renamed fixtures."""
name = updated_node.value
if name in self.renames:
new_name = self.renames[name]
return updated_node.with_changes(value=new_name)
return updated_node
Matching fixtures
def test_user():
...
FunctionDef(
name=Name(value='test_user'),
params=Parameters(),
body=SimpleStatementSuite(body=[
Expr(value=Ellipsis()),
]),
decorators=[],
)
Matching fixtures
@pytest.fixture
def test_user():
...
FunctionDef(
name=Name(value='test_user'),
params=Parameters(),
body=SimpleStatementSuite(body=[
Expr(value=Ellipsis()),
]),
decorators=[
Decorator(
decorator=Attribute(
value=Name(value='pytest'),
attr=Name(value='fixture'),
),
),
],
)
Matching fixtures
def is_pytest_fixture(node: cst.FunctionDef) -> bool:
for decorator in node.decorators:
match decorator.decorator:
# Handle @fixture
case cst.Name(value="fixture"):
return True
# Handle @pytest.fixture
case cst.Attribute(value=cst.Name(value="pytest"),
attr=cst.Name(value="fixture")):
return True
return False
class Transformer(cst.CSTTransformer):
def __init__(self):
self.renames: dict[str, str] = {}
def visit_FunctionDef(self, node) -> bool:
"""Collect fixtures that need to be renamed."""
if is_pytest_fixture(node) and should_rename(node):
old_name = node.name.value
self.renames[old_name] = generate_new_name(old_name)
return True # Continue visiting children
def leave_Name(self, original_node, updated_node):
"""Update variables that match renamed fixtures."""
name = updated_node.value
if name in self.renames:
new_name = self.renames[name]
return updated_node.with_changes(value=new_name)
return updated_node
Recap. That's basically all there is. And we can rename 100 or 1000 fixtures that are prefixed with
`test_`.
Running codemods
Clean git working tree
Run codemod script
Run formatters/linters
Commit just this
Automated changes = isolated commits
So this is the stuff I talk about at parties.
Workflow to run codemods!
1. Make sure git status says everything is clean. You don't want unstaged changes
2. Run the codemod. If something went wrong, just collect a couple examples and git reset --hard
3. Run formatters/linters (I mention this as a separate step)
Run it as part of your script, or I like precommit for that
4. Commit _just that_
This is a bit of boring advice but I'd feel bad not mentioning it.
The rule:
This is important! You will thank yourself if you need to rebase on top of a lot of changes.
Instead of dealing with conflicts it's often easier to drop the automated changes and re-run the script.
Writing codemods
Maintainability? No
90%? Claim success
Leave formatting to tooling: more fun
Test Driven Development
So advice on running codemods: be rigorous. Follow the method, atomic commits... BORING.
Now what do we care about when we're _writing_ codemods?
**PRESS**
this doesn't always apply, if you're a library author writing a codemod so all your users can upgrade,
you'll need a higher bar.
You probably can't assume your users all have a formatter like black or ruff.
**PRESS**
if the script gets you 90% of the way there, that might be enough for your use case!
It's fine to do things manually too, it's just another tool in the toolbox.
As long as you keep automated changes and manual changes in separate git commits.
I prefer to give control of formatting to a tool like black or ruff. Then the refactor script can mess
everything up - there's no need to manage whitespaces so the code still looks good after applying the
script.
Similar for things like duplicate or unused imports. If you have tooling doing that, your refactoring
scripts can add 3 times the same import.
Unix philosophy of combining small tools that do one thing well.
All in all: If you're targeting one codebase, for a one-off change: that's a bit liberating :).
Also, you might notice these bullet points are really where AI Agents strive.
Example: limitations (many)
Same-named local variables
Defined "after use"
Fixtures across files
Matching not the most robust (ex: pytest alias)
If we go back to our example, I followed my own advice about 'claiming success' at 90%.
To be fair: it was an actual success on the codebase I did that. And that's all it takes.
But let's look at where it falls short.
"after use" in quotes because it's really "after use" in the text or in the tree.
Not in the execution.
So many limitations to this script.
BUT SOMETIMES. ITS GOOD ENOUGH!
Example: Defined "after use"
before.py
def test_function(test_user):
assert test_user["name"] == "test"
@pytest.fixture
def test_user():
return {"name": "test"}
after.py
def test_function(user_fixture):
assert user_fixture["name"] == "test"
@pytest.fixture
def user_fixture():
return {"name": "test"}
Example: same-name variables
before.py
@pytest.fixture
def test_user():
return {"name": "test"}
def test_another_function():
test_user = {"name": "local"}
assert test_user["name"] == "local"
after.py
@pytest.fixture
def user_fixture():
return {"name": "test"}
def test_another_function():
test_user = {"name": "local"}
assert test_user["name"] == "local"
Going further with LibCST
State across files: multiple passes [Metadata APIs]
Variable scope management [ScopeProvider]
Pattern-matching nodes [matchers]
When needed LibCST gives you tools to go further.
But again, I want to stress that what works for you, is good enough.
Pattern-matching nodes: matchers are convenient to avoid a soup of if statements.
The `match` python statement can also help.
More robust matching in the previous slide: this could use the QualifiedName metadata provider. Would
simplify it.
AI will default to the if soup.
Automated refactoring: getting started
First step: use existing tools!
django-upgrade, pyupgrade
npx @next/codemod upgrade canary
Start asking "would this be doable?"
This is the main point I'm trying to make so I want to reinforce that.
Understanding the concept is great.
Look for existing tools! If you know it's possible, you know something can exist.
So maybe you find django-upgrade and you're set.
Would this be doable? or Why am I even doing this manually?
Defer all repetitive tasks to your enthusiastic coworker. They can write a script!
It's more fun. And it helps understanding the tools we use better, since syntax trees are at the heart
of a lot of devtools.
Ideas of refactoring scripts
Unittest to pytest
before.py
class TestAssertNotEqual(TestCase):
def test_you(self):
self.assertNotEqual(abc, 'xxx')
after.py
def test_you(self):
assert abc != 'xxx'
From https://github.com/pytest-dev/unittest2pytest/tree/main
This is a test from a popular repository that did not use libcst. The codemod is not compatible with recent python versions. It's a good candidate for a libcst port.
Ideas of refactoring scripts
Rewrite apps.get_model (Django) to local imports
before.py
Message = apps.get_model("Chat", "Message")
after.py
from chat.models import Message
The elephant
Time to address it. It's in the room. The elephant is AI of course.
Why would you write a LibCST script when you can just ask AI to "change this, do that"?
LibCST vs the Claudes
Deterministic vs faith-based
Easy rebasing vs another 30 minutes of GPU time
Claude can write the LibCST transformer
faith-based that you don't get fancy-creative Claude for some files.
I'll say it used to not work well because it consumed too much context.
But today most tooling will use 'sub agents' automatically, split the work between them.
So context is not bloated and they don't become end-of-week Claude.
Before: Claude can write the LibCST transformer.
ALSO... IF YOU INSIST...
LibCST vs the Claudes
Clear and obvious conclusion:
AI really good in many cases :)
Codemods feel better for some changes
Large codebases, reusability: push towards codemods
The Abstract Syntax Tree is one of these things that might sound mysterious the first time you hear it.
But when you look at it, it's actually quite approachable, and gives you a lot of power for things like
automated refactoring.
You won't necessarily use it every day. But knowing it powers some of your tools, and it's there if you
need it. Next time you have a very boring migration to make on your entire codebase, you might be able
to write a migration tool to automate part of
it.
Or find one on the internet :).
Resources
for reference in the slides
Thank you
Come talk to me!
If you're an expert about this. If you thought this was boring. If you were playing candy crush.
Syntax Trees Everywhere
Interpreter
Linters
Refactoring
'Transpiling': converting from one language to another
Template engines
I want to mention this as a motivation. It's nice to learn about syntax trees because they show up in so
many developer tools.
So: where else are ASTs used.
Interpreter: If you access the JavaScript AST and run it in Python, you have a JavaScript interpreter
written in Python.
You can also do a Python interpreter written in Python. That's one way of confusing people.
Linters: eslint in JavaScript, flake8 in Python.
formatters too: black uses the AST.
Transpiling: babel is an example in JavaScript to convert new syntax to older, compatible syntax. So we
can use the fancy new JavaScript features, and it converts the code back to an equivalent version taht
all browsers can understand.
Template engines: jinja2 is one example.
I think it's a fundamental concept, not because it is basic or simple - it's not - but because it powers
a lot of the tools we use.
And while it might be hard to write these tools - handling edge cases, different python versions... - ,
understanding the principles goes a long way.
Big picture: from source code to execution
Before we focus on the AST, let's talk about a very high-level view of how Python code gets executed.
This is so that you have an idea of the big picture.
You can look into it if you're interested.
PRESS
Python starts by reading your code as text, presumably from a file.
Big picture: from source code to execution
Tokenizing: splitting the text in relevant words.
Parsing: taking these tokens and grouping them according to the Python grammar.
From there, construction of the Abstract Syntax Tree.
There is a parse tree intermediate step, but we're skipping it.
Big picture: from source code to execution
Then the AST is compiled to bytecode instructions.
Big picture: from source code to execution
We are going to focus on the AST part.
DONT WORRY TOO MUCH THOUGH if this is new to you and you're a bit confused here.
You can understand and use the Abstract Syntax Tree without looking at the other parts.
There are many big words on this slide.
