计算python程序的数据依赖图

Question

我有一段python代码：a0 = 1a1 = 2a2 = add(a0, a1)a3 = [i + 1 for i in range(a2)]我想得到以下数据依赖图：{ "a0": [], "a1": [], "a2": ["a0", "a1"], "a3": ["a2"]}我可以locals()用来获取局部变量的列表，但这就是我现在所拥有的。我可以尝试解析 AST，但它似乎是一个兔子洞，因为我需要处理 AST 的每一个怪癖。有任何想法吗 ？用例：我正在尝试编写一种小型 DSL 语言，其中变量是任务，代表需要运行的长时间运行的命令。我正在使用doit但我发现在代码中已经可见时必须显式声明依赖项很麻烦且难以阅读。

Answer 1

最后ast模块已经有了很方便的“walk”方法。

import ast

def get_deps(code):

    body = ast.parse(code)

    _, statements = next(ast.iter_fields(body))

    full_graph = {

        assign.targets[0].id: [

            d.id for d in ast.walk(assign) if isinstance(d, ast.Name)

        ]

        for assign in statements

    }

    # full_graph also contains `range` and `i`. Keep only top levels var

    restricted = {}

    for var in full_graph:

        restricted[var] = [d for d in full_graph[var] if d in full_graph and d != var]

    return restricted

if __name__ == "__main__":

    d = get_deps(

        """

a0 = 1

a1 = 2

a2 = add(a0, a1)

a3 = [i + 1 for i in range(a2)]

"""

    )

    assert d == {"a0": [], "a1": [], "a2": ["a0", "a1"], "a3": ["a2"]}, d