usm_draw.py 文件源码

def generate_graph(usm):
    current = usm.get_root()
    G = nx.DiGraph()
    labels = {}
    idx = 0
    queue = deque([(idx, current)])
    colors = [_get_color(current)]
    while len(queue) != 0:
        pidx, current = queue.popleft()
        for key, child in current.children.items():
            idx+=1
            G.add_node(idx, tree_node=child)
            colors.append(_get_color(child))
            G.add_edge(pidx, idx, label=key)
            labels[idx] = key
            queue.append((idx, child))
    return G, labels, colors