Edgelist to Tree

trees.py

if __name__ == "__main__":

    import pprint
    import copy

    engine = init_engine(DB_CONN)

    children = [(1,2),(1,3),(2,4)]

    edgelist = engine.execute(children)
    
    tree = {}
    
    seen_nodes = set()

    def addLeaf(tree, parent, child):
        signal = False

        if parent in tree:
            tree[parent].update({child: {}})
            signal = tree
        
        elif tree:
            for branch in tree.values():
                if addLeaf(branch, parent, child):
                    signal = True
        
        else:
            signal = False

        return signal

    for parent, child in edgelist:
        
        if child in seen_nodes:
            tree[parent].update({child: copy.deepcopy(tree[child])})
            del tree[child]
        else:
            result = addLeaf(tree, parent, child)
        
            if not result:
                tree[parent] = {child: {}}
        
        seen_nodes.update([parent, child])
        

    pprint.pprint(tree)