From 1dbe0975fcc78b4518aa3742528c6c764f281f65 Mon Sep 17 00:00:00 2001
From: Stephen Diehl <stephen.m.diehl@gmail.com>
Date: Mon, 9 Apr 2012 07:55:11 -0400
Subject: [PATCH] NumpyChannel

---
 zipline/zmq_utils.py | 60 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/zipline/zmq_utils.py b/zipline/zmq_utils.py
index e32931dc..5e0a529f 100644
--- a/zipline/zmq_utils.py
+++ b/zipline/zmq_utils.py
@@ -3,6 +3,9 @@ Misc ZeroMQ utilities.
 """
 import gevent
 import msgpack
+import numpy
+from numpy import dtype
+from pandas import DataFrame
 from gevent_zeromq import zmq
 
 from contextlib import closing
@@ -92,3 +95,60 @@ def ZmqConsole(sock_typ, socket_addr, sock_conn=None, context=None):
             import pdb; pdb.set_trace()
 
     return gevent.spawn(console)
+
+class NumpyChannel(zmq.Socket):
+
+    def recv_pandas(self, flags=0, copy=True, track=False):
+
+        # Pandas Metadata
+        index, columns, dtype_name, shape = msgpack.loads(self.recv(flags=flags))
+
+        # Pandas ndarray
+        ndbuffer = self.recv(flags=flags, copy=copy, track=track)
+        buf = buffer(ndbuffer)
+
+        ndarray = numpy.frombuffer(buf, dtype=dtype(dtype_name)).reshape(shape)
+        return DataFrame(data=ndarray, index=index,
+                columns=columns, dtype=dtype_name)
+
+    def send_pandas(self, df, flags=0, copy=True, track=False):
+
+        # Pandas Metadata
+        index = df.index.tolist()
+        columns = df.columns.tolist()
+        dtype_name = df.values.dtype.name
+        shape = df.values.shape
+
+        # Pandas ndarray
+        ndarray = df.values
+
+        metadata = msgpack.dumps((index, columns, dtype_name, shape))
+
+        self.send(metadata, flags|zmq.SNDMORE)
+        return self.send(ndarray, flags, copy=copy, track=track)
+
+if __name__ == '__main__':
+
+    from numpy.random import randn
+    df = DataFrame(randn(5,5))
+
+    ctx = zmq.Context.instance()
+
+    def send():
+        pub = NumpyChannel(ctx, zmq.PUSH)
+        pub.bind('inproc://a')
+
+        for i in xrange(100):
+            pub.send_pandas(df, copy=False)
+
+    def recv():
+        sub = NumpyChannel(ctx, zmq.PULL)
+        sub.connect('inproc://a')
+
+        for i in xrange(100):
+            sub.recv_pandas(copy=False)
+
+    gevent.joinall([
+        gevent.spawn(send),
+        gevent.spawn(recv)
+    ])