Source

asyncplot / asyncplot / server.py

Full commit
r"""Threaded Socket Server"""

import cPickle
import socket
import subprocess
import threading
import SocketServer

import mmf


class PlotRequestHandler(SocketServer.BaseRequestHandler):
    def handle(self):
        self.request.sendall(self.server.get_data())
        #self.request.shutdown(socket.SHUT_RDWR)
        #self.request.close()


class PlotSocketServer(SocketServer.ThreadingMixIn,
                       SocketServer.TCPServer):
    r"""Simple threaded socket server.  Listens on the specified port, and
    launches Socket threads on connection.

    Call :meth:`update` to update the plot data.

    To Do
    -----
    * Add timeouts.
    """
    _PORT = 9999
    _HOST = ''

    def __init__(self, host=_HOST, port=_PORT):
        self.lock = threading.Lock()  # RLock()??
        self.have_data = threading.Event()
        self.data = None
        
        self.daemon_threads = True       # Don't keep python alive.
        self.allow_reuse_address = True  # Allow port reuse after termination.
        SocketServer.TCPServer.__init__(self, (host, port),
                                        PlotRequestHandler)

    def get_data(self, block=True):
        r"""The object is pickled here."""
        self.have_data.wait()
        with self.lock:
            return cPickle.dumps(self.data)
        
    def update(self, *v, **kw):
        with self.lock:
            self.data = (v, kw)
        self.have_data.set()    # This is very slow!


class Server(threading.Thread):
    r"""If you pass `spawn_client_command`, then a plot client will be spawned
    in a subprocess.
    """
    def __init__(self,
                 host=PlotSocketServer._HOST,
                 port=PlotSocketServer._PORT,
                 client_cmd=None,
                 wait_for_client=False):
        r"""
        Parameters
        ----------
        client_cmd : str
           If provided, then this command will be spawned in a subprocess (as a
           list of commands + arguments expected by :func:`subprocess.call`).
        wait_for_client : Bool
           If `True`, then pause server to allow the user to adjust plot
           windows.
        """
        self.client_cmd = client_cmd
        self.wait_for_client = wait_for_client
        self.server = PlotSocketServer(host=host, port=port)
        threading.Thread.__init__(self, target=self.server.serve_forever)
        self.daemon = True               # Don't keep python alive.

    def __enter__(self):
        self.start()
        if self.client_cmd is not None:
            self.client_process = subprocess.Popen(self.client_cmd)
        else:
            self.client_process = None
        
        if self.wait_for_client:
            raw_input("Adjust plot windows and press return to continue...")

        return self.server

    def __exit__(self, type, value, traceback):
        self.server.shutdown()      # This does not close the socket
        if self.client_process is not None:
            self.client_process.terminate()
        self.server.server_close()  # This actually closes the socket
        return value is None


def example_server():
    import time
    import numpy as np
    x = np.linspace(0, 2*np.pi, 100)
    t = 0
    dt = 0.01
    tic = time.time()
    with Server(client_cmd=['ipython', '--pylab=osx', 'client.py'],
                wait_for_client=True) as s:
        for n in xrange(1000000):
            if n % 100000 == 0:
                print n, time.time() - tic
                tic = time.time()
            t += dt
            y = np.sin(x + t)
            s.update(x, y)