I'm unfamiliar with the AST module and would appreciate any insight. If, for example, I have a string that contains a valid python script such as
import sys #Just any module
class SomeClass:
def __init__(self):
self.x = 10
self.b = 15
def a_func(self):
print(self.x)
I would like to be able to programmatically edit lines such as changing self.x = 10 to something like self.x = 20. I can break it down somewhat with ast via:
some_string = "..." #String of class above
for body_item in ast.parse(some_string):
...
But this doesn't feel like the "right" way(not that there is a right way since this is somewhat niche). I was hoping someone could correct me towards something cleaner, or just better.
You can start by using ast.dump to get an idea of the AST structure of the code you're dealing with:
import ast
code='self.x = 10'
print(ast.dump(ast.parse(code), indent=2))
This outputs:
Module(
body=[
Assign(
targets=[
Attribute(
value=Name(id='self', ctx=Load()),
attr='x',
ctx=Store())],
value=Constant(value=10))],
type_ignores=[])
From which you can see what you want to look for is an Assign node where the first of targets is an Attribute node whose value is a Name node with an id of 'self' and an attr of 'x'.
With this knowledge, you can then use ast.walk to traverse the AST nodes to look for a node with the aforementioned properties, modify its value to a Constant node with a value of 20, and finally use ast.unparse to convert AST back to a string of code:
import ast
code = '''
import sys #Just any module
class SomeClass:
def __init__(self):
self.x = 10
self.b = 15
def a_func(self):
print(self.x)
'''
tree = ast.parse(code)
for node in ast.walk(tree):
if (
isinstance(node, ast.Assign) and
isinstance((target := node.targets[0]), ast.Attribute) and
isinstance(target.value, ast.Name) and
target.value.id == 'self' and
target.attr == 'x'
):
node.value = ast.Constant(value=20)
print(ast.unparse(tree))
This outputs:
class SomeClass:
def __init__(self):
self.x = 20
self.b = 15
def a_func(self):
print(self.x)
Note that ast.unparse requires Python 3.10 or later. If you're using an earlier version, you can use astunparse.unparse from the astunparse package instead.