Python ast to dot graph - PullRequest
Новая
       29

Python ast to dot graph

8 голосов
/ 01 декабря 2011

Я анализирую AST, сгенерированный кодом Python, для «удовольствия и прибыли», и мне хотелось бы иметь что-то более графическое, чем «ast.dump», чтобы фактически увидеть сгенерированный AST.это уже дерево, поэтому создание графика не должно быть слишком сложным, но я не понимаю, как я мог бы это сделать.

ast.walk, кажется, подходит для стратегии BFS, и visitXметоды, которые я не вижу на самом деле, или я, кажется, не могу найти способ создать граф ...

Кажется, что единственный способ - написать свою собственную функцию обхода DFS, разве это не так?имеет смысл?

Ответы [ 2 ]

6 голосов
/ 01 декабря 2011

Фантастика, это работает, и это действительно просто 

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)
6 голосов
/ 01 декабря 2011

Если вы посмотрите на ast.NodeVisitor, это довольно тривиальный класс. Вы можете либо подклассить его, либо просто переопределить его стратегию ходьбы на то, что вам нужно. Например, сохранить ссылки на родительский элемент при посещении узлов очень просто для реализации этого способа, просто добавьте метод visit, который также принимает родительский элемент в качестве аргумента, и передайте его из собственного generic_visit.

P.S. Кстати, похоже, что NodeVisitor.generic_visit реализует DFS, поэтому все, что вам нужно сделать, это добавить передачу родительского узла.

...