Python ast 转换为点图

Question

我正在为“有趣和收益”的目的分析由python代码生成的AST，并且希望能够拥有比“ast.dump”更加图形化的工具来实际查看生成的AST。

理论上说它已经是一棵树了，所以创建一个图应该不会太难，但是我不知道该如何做。

ast.walk似乎使用BFS策略进行遍历，而visitX方法中我无法看到父节点或找到创建图的方法...

似乎唯一的办法是编写自己的DFS遍历函数，这种做法有意义吗？

- andrea_crotti

Answer 1

太棒了，它可以正常工作，而且非常简单。

class AstGraphGenerator(object):

    def __init__(self):
        self.graph = defaultdict(lambda: [])

    def __str__(self):
        return str(self.graph)

    def visit(self, node):
        """Visit a node."""
        method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node)

    def generic_visit(self, node):
        """Called if no explicit visitor function exists for a node."""
        for _, value in ast.iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, ast.AST):
                        self.visit(item)

            elif isinstance(value, ast.AST):
                self.graph[type(node)].append(type(value))
                self.visit(value)

这个问题与普通的NodeVisitor一样，但是我使用了一个defaultdict，用于为每个子节点添加节点类型。然后，我将这个字典传递给pygraphviz.AGraph，并获得了漂亮的结果。

唯一的问题是类型信息并不太具体，但是使用ast.dump()则过于冗长。

最好的方法是获取每个节点的实际源代码，这可能吗？

编辑：现在情况好多了，我在构造函数中也传入了源代码，并尝试获取代码行，如果不可能，则只打印出类型。

class AstGraphGenerator(object):

    def __init__(self, source):
        self.graph = defaultdict(lambda: [])
        self.source = source  # lines of the source code

    def __str__(self):
        return str(self.graph)

    def _getid(self, node):
        try:
            lineno = node.lineno - 1
            return "%s: %s" % (type(node), self.source[lineno].strip())

        except AttributeError:
            return type(node)

    def visit(self, node):
        """Visit a node."""
        method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node)

    def generic_visit(self, node):
        """Called if no explicit visitor function exists for a node."""
        for _, value in ast.iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, ast.AST):
                        self.visit(item)

            elif isinstance(value, ast.AST):
                node_source = self._getid(node)
                value_source = self._getid(value)
                self.graph[node_source].append(value_source)
                # self.graph[type(node)].append(type(value))
                self.visit(value)

- andrea_crotti

Answer 2

如果你看一下ast.NodeVisitor，它是一个相当简单的类。你可以通过继承它或重新实现其遍历策略来满足你的需求。例如，在访问节点时保持对父节点的引用非常容易实现，只需添加一个接受父节点作为参数的visit方法，并从你自己的generic_visit中传递它。

顺便说一下，看起来NodeVisitor.generic_visit实现了DFS，所以你只需要添加父节点传递即可。

- Eli Bendersky

---

是的，你说得对，这是一个非常简单的实现。起初我认为我必须检查所有可能的情况，但实际上只需要检查它是否是一个列表就足以遍历整棵树了。非常感谢。 - andrea_crotti