如何覆盖python脚本文件中的一些赋值？

Question

我有一个描述一组对象/变量的python文件。使用file.py的示例

a = 5
b = a+1

我想导入文件的内容，并且能够覆盖某些变量的值：parse_with_overwrite(file.py, {"a": 6})将使当前的locals()包含a=6和b=7。

当然，与任何python文件一样，并非所有变量都在一行…中描述所以不可能在=符号上拆分每一行。此外，变量的顺序很重要，因为一些变量可能依赖于其他变量。

有没有什么内置的或者库可以帮助我做到这一点？

目前，我通过以下方式获取file.py的内容：

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

Answer 1

一种安全的方法是使用ast module来解析文件。然后，您可以修改ast中的节点，最后对其进行编译。

这是一个开始：

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)