AharonSambol / PrettyPrintTree

Python library to print trees

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Pretty print non-class structures

zademn opened this issue · comments

It seems that when I try to print a non-class structure (such as a collection) it defaults to json (Or maybe I'm misusing something). Suppose my nodes are defined by the tuple (lvl, idx) where (0, 0) is the root

pt = PrettyPrintTree(
    lambda lvl, idx: nodes_below((lvl, idx), lvl + 1) if lvl < 3 else [],
    lambda x: x
)

pt((0, 0)) 
# returns
  JSON 
 ┌─┴─┐
 0   0 

# And I expect a depth 3 binary tree, with `(lvl, idx)` as nodes

Where nodes_below is a general function that gives the nodes below a node, given a level. For example, for (0, 0), lvl = 1 it will return [(1, 0), (1, 1)], for (1, 1), lvl = 2 it will return [(2, 2,) (2, 3)] etc.

def nodes_below(node: tuple[int, int], level: int) -> list[tuple[int, int]]:
    node_lvl, node_idx = node
    levels_below = level - node_lvl
    start_idx = node_idx * (2 ** levels_below) # Binary tree
    nodes = [(level, idx) for idx in range(start_idx, start_idx + (2 ** levels_below))]
    return nodes

Fixed!