Print graphical lambda calculus

print_lambda.py

from arpeggio.cleanpeg import ParserPEG
from arpeggio import PTNodeVisitor, visit_parse_tree

GRAMMAR=r"""
full = expr EOF
expr = var / function / application
function = "/" var "." expr
application = "(" expr expr ")"
var = r"\w+"
"""

parser = ParserPEG(GRAMMAR, root_rule_name="full")

def _make_top_pipe(cols):
    last_non_pipe=max(i for i, col in enumerate(cols) if col[0]!=" ")
    for i in range(last_non_pipe+1, len(cols)):
        cols[i]="─"+cols[i][1:]
    if cols[last_non_pipe][0]=="┐":
        cols[last_non_pipe]="┬"+cols[last_non_pipe][1:]
    return cols


class MyVisitor(PTNodeVisitor):
    def visit_full(self, node, children):
        return children[0]
    
    def visit_expr(self, node, children):
        return children[0]
    
    def visit_function(self, node, children):
        var, expr = children
        return _make_top_pipe(expr) + ["─"] + var
        
    def visit_application(self, node, children):
        func, arg = children
        return _make_top_pipe(func) + ["─"] + [("┐" if i==0 else " ") + col for i, col in enumerate(arg)]
    
    def visit_var(self, node, children):
        return list(str(node))

def p(text):
    tree = parser.parse(text)
    tree2 = visit_parse_tree(tree, MyVisitor())
    print("\n".join("".join(row) for row in zip_longest(*tree2, fillvalue=" ")))

for name, lam in [
    ("Abstraction", "/x.M"),
    ("Application", "(f x)"),
    ("true", "/x./y.x"),
    ("false", "/x./y.y"),
    ("and", "/a./b.((a b) a)"),
    ("0", "/f./x.x"),
    ("1", "/f./x.(f x)"),
    ("2", "/f./x.(f (f x))"),
    ("succ", "/n./f./x.(f ((n f) x))"),
    ("Y", "/g.(/x.(g (x x)) /x.(g (x x)))"),
    ("Omega", "(/x.(x x) /x.(x x))"),
    
]:
    print(f"{name}: {lam}")
    p(lam)
    print()