Python ast 转点图

Question

我正在分析由 python 代码生成的 AST，以获取“乐趣和利润”，我希望有比“ast.dump”更图形化的东西来实际查看生成的 AST。

理论上已经是一棵树，所以创建一个图应该不会太难，但我不明白我该怎么做。

ast.walk 似乎以 BFS 策略行走，而 visitX 方法我无法真正看到父级，或者我似乎没有找到创建图形的方法...

看来唯一的办法就是自己写一个 DFS walk 函数，这有意义吗？

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)