Draw out Black AST nodes

explain.py

#!/usr/bin/env python3
from black import lib2to3_parse
from blib2to3.pytree import Leaf, type_repr

func_example = """
def f(x, y):
   # add two operands

   return (x + y)  # yes just this
"""

parsed_func = lib2to3_parse(func_example)


def es_dot(s):
    return s.replace('"', '\\"')  # .replace("\n", r"\l")


def fmt_html(s):
    return html.escape(s).replace("\n", "<br/>")


import html
from blib2to3.pgen2.token import tok_name


def node_xml(node, node_id):

    node_name = es_dot(str(node))

    if isinstance(node, Leaf):
        # we now have a leaf description
        contents = f"<tr><td>Content</td><td>{fmt_html(node.value)}</td></tr>"
        leaf_descr = f"<tr><td>Prefix</td><td>{fmt_html(node._prefix)}</td></tr>"
        node_type_descr = f"<tr><td>Type</td><td>{tok_name[node.type]}</td></tr>"
    else:
        contents = f"""<tr><td>{fmt_html(node_name)}</td></tr>"""
        leaf_descr = ""
        node_type_descr = ""
    # <tr><td>{html.escape(str(type(node)))}</td></tr>
    return f"""
    {node_id}[label=<
    <table border="0" cellborder="1" cellspacing="0" align="left">
        {node_type_descr}
        {contents}
        {leaf_descr}
    </table>
    >]
    """


def to_digraph(node):
    out = """
    digraph g {
    node[shape=plaintext];
    """
    nodes = [node]
    ids = {}

    def get_id(node):
        if id(node) not in ids:
            ids[id(node)] = f"node_{len(ids)}"
        return ids[id(node)]

    while nodes:
        node = nodes.pop()
        node_name = es_dot(str(node))
        node_id = get_id(node)
        node_type = tok_name.get(node.type, "(unknown)")
        # let's add the node here
        out += node_xml(node, node_id)
        if node.parent:
            out += f"""
            {get_id(node)} -> {get_id(node.parent)}
            """
        nodes += node.children
    out += """
    }
    """

    return out


print(to_digraph(parsed_func))