Skip to content

Instantly share code, notes, and snippets.

@rtpg
Created July 5, 2021 09:53
Show Gist options
  • Save rtpg/b687babfd1befaab719f9b89add70ad1 to your computer and use it in GitHub Desktop.
Save rtpg/b687babfd1befaab719f9b89add70ad1 to your computer and use it in GitHub Desktop.
Draw out Black AST nodes
#!/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))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment