Question

我有一段python代码：

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

我想获得以下依赖关系图：

{ "a0": [], "a1": [], "a2": ["a0", "a1"], "a3": ["a2"]}

我可以使用locals()来获取局部变量的列表，但这只是我现在所拥有的。我可以尝试解析AST，但这似乎是个小难题，因为我需要处理AST的所有问题。有什么想法吗？

用例：我正在尝试编写一种小的DSL语言，其中的变量是task，代表需要运行的长期运行的命令。我正在使用doit，但是我发现麻烦的是，必须在代码中已经显式声明依赖项的情况下，才必须显式声明依赖项。

Answer 1

最后，ast模块已经具有非常方便的“步行”方法。

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

计算一段python代码的依赖关系图

1 个答案: