python - 计算Python程序的数据依赖图

Question

我有一段Python代码:

a0 = 1
a1 = 2
a2 = add(a0, a1)
a3 = [i + 1 for i in range(a2)]

我想得到以下data-dependency graph :

{ "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