Python ast to dot graph

Question

Estoy analizando el AST generado por el código python para "diversión y beneficio", y me gustaría tener algo más gráfico que "ast.dump" para ver realmente el AST generado.

En teoría ya es un árbol, por lo que no debería ser demasiado difícil crear un gráfico, pero no entiendo cómo podría hacerlo.

ast.walk parece caminar con una estrategia de BFS, y los métodos visitX que realmente no puede ver el padre o yo no parecen encontrar una manera de crear un gráfico ...

Parece como que la única forma es escribir mi propia función de caminar DFS, ¿tiene sentido?

Answer 1

Si mira ast.NodeVisitor, es una clase bastante trivial. Puede subclasificarlo o simplemente volver a implementar su estrategia de caminar a lo que necesite. Por ejemplo, mantener las referencias al padre cuando se visitan los nodos es muy simple de implementar de esta manera, simplemente agregue un método visit que también acepte el padre como argumento, y páselo desde su propio generic_visit.

P.S. Por cierto, parece que NodeVisitor.generic_visit implementa DFS, por lo que todo lo que tiene que hacer es agregar el nodo primario que pasa.

Answer 2

fantástico, funciona y es realmente sencilla

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) 

Por lo que es lo mismo que un NodeVisitor normal, pero tengo una defaultdict donde agrego el tipo de nodo para cada hijo. Luego paso este diccionario a pygraphviz.AGraph y obtengo mi buen resultado.

El único problema es que el tipo no dice mucho, pero por otro lado el uso de ast.dump() es demasiado detallado.

Lo mejor sería obtener el código fuente real para cada nodo, ¿es posible?

EDITAR: ahora es mucho mejor, paso en el constructor también el código fuente e intento obtener la línea de código si es posible, de lo contrario solo imprimo el tipo.

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)