How to overwrite some assignments in a python script file?

Question

I have a python file describing a set of objects/variables. Example with file.py :

a = 5
b = a+1

And I would like to import the file's content with the ability to overwrite the value of some variables: parse_with_overwrite(file.py, {"a": 6}) would make the current locals() contain a=6 and b=7 .

Of course, like any python file, not all variables are described in one single line… so it's not possible to split each line on the = sign. Moreover, order of the variables matters since some variables may depend on other.

Is there any builtin or library that could help me achieve this?

Currently, I get the content of file.py with:

context = {}
with open("file.py", "r") as f:
    # the magic should occur here :-)
    exec(f.read(), None, context)

Answer 1

A safe way to do it is to use the ast module to parse the file. You can then modify the nodes in the ast and finally compile it.

Here is a start:

import ast

def replace_top_assignments(tree, map):
    for i, node in enumerate(tree.body):
        if not isinstance(node, ast.Assign):
            continue

        for nt in node.targets:
            if not isinstance(nt, ast.Name):
                # deal with tuple assignments a, b = 1, 2
                continue
            if nt.id in map:
                # ok we have to replace this
                if not isinstance(node.value, ast.Constant):
                    # warning it is not assigned a constant
                    continue
                # change value for all targets. this will break a=b=2 with a replace of a = 5
                node.value.value = map[nt.id]

src = """
a = 23
b = a
# todo handle these
# b = c = 2
# r, s = 1, a
# conf = usion = 1  what should happen here if you replace usion with 4 ? is conf 4?

print(a)
print(b)
"""

tree = ast.parse(src)
replace_top_assignments(tree, {'a': 4})
code_obj = compile(tree, 'aaa', mode='exec')
exec(code_obj)