abp/abp/fancy.py

import time, atexit, json
import sys
import networkx as nx
import numpy as np
import websocket
from socket import error as socket_error
import graphstate
import clifford
import util

class GraphState(graphstate.GraphState, nx.Graph):
    def __init__(self, *args, **kwargs):
        graphstate.GraphState.__init__(self, *args, **kwargs)
        self.connect_to_server()

    def connect_to_server(self, uri = "ws://localhost:5000"):
        """ Attempt to connect to the websocket server """
        try:
            self.ws = websocket.create_connection(uri, timeout=0.1)
            atexit.register(self.shutdown)
        except: #TODO: bad practice
            self.ws = None

    def shutdown(self):
        """ Close the connection to the websocket """
        if not self.ws:
            return
        self.ws.close()

    def update(self, delay = 0.5):
        """ Call this function when you are ready to send data to the browser """
        if not self.ws:
            return

        # Automatically perform layout if position is not provided
        if not all(("position" in node) for node in self.node.values()):
            self.layout()

        # Send data to browser and rate-limit
        try:
            self.ws.send(json.dumps(self.to_json(stringify=True)))
            self.ws.recv()
            time.sleep(delay)
        except websocket._exceptions.WebSocketTimeoutException:
            print "Timed out ... you might be pushing a bit hard"
            time.sleep(delay)
            #self.ws.close()
            #self.connect_to_server()

    def layout(self):
        """ Automatically lay out the graph """
        pos = nx.spring_layout(self, dim=3, scale=np.sqrt(self.order()))
        middle = np.average(pos.values(), axis=0)
        pos = {key: value - middle for key, value in pos.items()}
        for key, (x, y, z) in pos.items():
            self.node[key]["position"] = util.xyz(x, y, z)

    def add_vops(self):
        """ Automatically add vops if they're not present """
        for key in self.node:
            if not "vop" in self.node[key]:
                self.node[key]["vop"] = clifford.identity