BasicNode.py

from ColorBank import ColorBank

class BasicNode:
    default_radius = 25
    default_color = (0,0,0)
    default_data = None
    default_position = (0, 0)

    def __init__(self, position=None,
                       radius=None,
                       color=None,
                       data=None,
                       randomize=None):
        if randomize is not None:
            r=randomize
            if data         is None: data       = '(random)'
            if color        is None: color      = ColorBank.random(r)
            if radius       is None: radius     = r.randint(3,50)
            if position     is None: position   = (r.random(), r.random())
        else:
            if data         is None: data       = BasicNode.default_data
            if color        is None: color      = BasicNode.default_color
            if radius       is None: radius     = BasicNode.default_radius
            if position     is None: position   = BasicNode.default_position

        if any(map(lambda c: not (0 <= c <= 1), position)):
            raise Exception('Woah there buddy.')

        self.data       = data
        self.color      = color
        self.radius     = radius
        self.position   = position
    def __str__(self):
        return str(self.data)
    def __repr__(self):
        return str(self.__dict__)

ColorBank.py

class ColorBank:
    def __init__(self):
        self.black = (0, 0, 0)
        self.white = (255, 255, 255)
        self.red   = (255, 0, 0)
        self.green = (0, 255, 0)
        self.blue  = (0, 0, 255)

    def set_color(self, name, red, green, blue):
        setattr(self, str(name), (red, green, blue))

    @classmethod
    def get_inverse(cls, color, alpha=1):
        inverses = [255 - c for c in color] + [alpha]
        return tuple((channel for channel in inverses))

    @classmethod
    def random(cls, r):
        return tuple((r.randint(0, 255) for i in range(3)))

for-gui.py

ListeyTree = [1,[2,3,4],5,6,[7,8,[9,10],11,12,[13,14],15]]

import pygame
import networkx as nx

from ColorBank import ColorBank
from BasicNode import BasicNode

def construct_tree_from_lisp(lisp):
    """
   Creates a Tree object from a LISP-style list:
   
       [1, 2, [3, 4, 5], [6, 7], 8]
   """
    if isinstance(lisp, list):
        v = lisp[0]
        c = []
        for child in lisp[1:]:
            c.append(construct_tree_from_lisp(child))
        return Tree(value = v, children = c)
    else:
        return Tree(value=lisp)
 
class Tree:
    """
   Children is a list of subtrees
   """
    def __init__(self, children=list(), value=None):
        self.children = children
        self.value    = value
       
    def minimum (self):
        if self.children:
            minchildren = [self.value]
            for item in self.children:
                minchildren.append(item.minimum())
            return min(minchildren)
        else:
            return self.value
       
    def get_child (self, v):
        return filter(lambda t:t.value is v, self.children)[0]

    def to_networkx(self, graph=None):
        if graph is None:
            graph = nx.DiGraph()

        for child in self.children:
            graph.add_edge(self.value, child.value)
            child.to_networkx(graph)

        return graph

    def __int__(self):
        return self.value

    def __repr__(self):
        return str(self.value)
        return '{}: [{}]'.format(str(self.value), ', '.join(map(int, self.children)))
 
tree1 = [1,[2,3,4],5,[6,7,[8]]]
tree2 = [[1,2],3, [4,[5,6,7,[8,9]]]]
 
class Zipper:
    def __init__ (self, tree):
        self.tree = construct_tree_from_lisp (tree)
        self.path = [self.tree]
 
       
    def move_to_child (self, child):
        assert child in self.tree.children
        self.tree = child
       
    def Zip (self, child):
        self.path.append (self.tree)
        self.move_to_child (child)
       
    def UnZip (self):
        self.tree = self.path.pop()
       
    def Add (self, item):
        self.tree.children.append(item)
       
    def Remove (self, item):
        assert item in self.tree.children
        assert type(item) is not instance  
        self.tree.children.remove (item)
        return self.tree
       
    def Min (self):
        return self.tree.minimum()

    def to_networkx(self):
        G = self.tree.to_networkx()
        setattr(G, '_my_root', self.tree.value)

        # set all nodes to be blue with radius 5
        for node in G.nodes():
            G.node[node]['color'] = (0, 0, 255)
            G.node[node]['radius'] = 5
            G.node[node]['data'] = node

        # Set the root node to be black
        G.node[self.tree.value]['color'] = (0, 0, 0)

        return G


class Visualizer:
    def __init__(self, size=(640, 480), graph=nx.Graph(), edge_width = 2):
        """where `size` is a 2-tuple representing screen dimens"""

        self.screen = pygame.display.set_mode(size)

        self.colors = ColorBank()
        self.graph = graph
        self.edge_width = edge_width
        self.layout_algorithms = [getattr(nx, a) for a in dir(nx) if a.endswith('_layout')]
        # TODO sometimes crashes here; why?
        pygame.font.init()
        self.text_font = pygame.font.SysFont('monospace', 15)

    def do_layout(self, layout_algorithm=nx.spring_layout):
#        try:
        p = layout_algorithm(self.graph)
        #except:
        #    print('Layout algorithm `{!s}` not yet supported.'.format(repr(layout_algorithm).split()[1]))
        #    print('Please install the appropriate package.')
        #    return

        for node, position in zip(p.keys(), p.values()): # in p isn't working: iteration over non-sequence
            self.graph.node[node]['position'] = ((position[0] + 1) / 2, (position[1] + 1) / 2)

    def draw(self):
        self.screen.fill(self.colors.white)
        size = self.screen.get_size()

        for src, dst in self.graph.edges():
            pygame.draw.line(self.screen, self.colors.green,
                             self.floats_to_pos(self.graph.node[src]['position']),
                             self.floats_to_pos(self.graph.node[dst]['position']), self.edge_width)

        for node, node_data in self.graph.nodes(data=True):
            normal_pos = self.floats_to_pos(node_data['position']) # keep track of z order for drag drop
            pygame.draw.circle(self.screen, node_data['color'], normal_pos, node_data['radius'], 0)
            label = self.text_font.render(str(node_data['data']), True, ColorBank.get_inverse(self.colors.white))
            self.screen.blit(label, normal_pos)

        pygame.display.update()

    def floats_to_pos(self, floats):
        return tuple((int(coordinate * scale) for coordinate, scale in zip(floats, self.screen.get_size())))

    def pos_to_floats(self, position):
        return tuple((coordinate / scale for coordinate, scale in zip(position, self.screen.get_size())))

    def loop(self, zipper, layout_func):
        """Runs the simulator.

        >>> pygame.init()
        (6, 0)
        >>> Visualizer(size=(640, 480), graph=make_graph()).loop()
        """
        ingame=True
        self.graph = zipper.to_networkx()
        self.do_layout(layout_algorithm=layout_func)

        while True:
            self.graph = zipper.to_networkx()
            self.do_layout(layout_algorithm=layout_func)
            self.draw()
            print 'did draw'
            pygame.time.delay(500)
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    ingame = False
        pygame.quit()

import generators
make_graph = lambda: \
    generators.random_graph((5, 20), .3,
                            data=(i for i in range(50)),
                            color=lambda r: ColorBank.random(r),
                            radius='int(3, 10)',
                            position=lambda r: tuple([r.random(), r.random()]))


def tree_layout(graph):
    nexts = [-.8] * 50

    ret = dict()

    scale_factor=8

    def minimum_ws(tree, depth=-6):
        ret[tree] = (nexts[depth], float(depth)/scale_factor)
        nexts[depth] += 1.0/scale_factor
        for c in graph.neighbors(tree):
            minimum_ws(c, depth + 1)

    minimum_ws(graph._my_root)

    return ret

v=Visualizer(size=(640, 480))

Zippie = Zipper(ListeyTree)
#print Zippie.tree.children
#print Zippie.tree.value
#print Zippie.path
# 
#print ""
#print ""
 
import thread
thread.start_new_thread(v.loop, (Zippie, tree_layout))

Zippie.Zip (Zippie.tree.get_child(2))

#print Zippie.path
Zippie.UnZip ()


#print ""
#print Zippie.path

#quit()