abp/abp/fancy.py

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

class GraphState(graphstate.GraphState, networkx.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)
            atexit.register(self.shutdown)
        except socket_error:
            self.ws = None

    def shutdown(self):
        """ Close the connection to the websocket """
        self.update()
        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
        self.ws.send(json.dumps(self.to_json()))
        time.sleep(delay)

    def layout(self, dim=3):
        """ Automatically lay out the graph """
        pos = networkx.spring_layout(self, dim, 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)