Edit attribute in script string with AST

Question

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.

Asked By: concernedDev

||

Source

Answer 1

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.

Demo: https://trinket.io/python3/3b09901326

Answered By: blhsing

Edit attribute in script string with AST

Question:

Answers: