Merge pull request #30 from quantopian/stable_upstream

Stable upstream

zipline/date_utils.py

@@ -0,0 +1,79 @@
+from collections import namedtuple
+
+import time
+import pytz
+import calendar
+from dateutil import rrule
+from datetime import datetime, date, timedelta
+from dateutil.relativedelta import *
+
+# Datetime Tuple
+# --------------
+d_tuple = namedtuple('dt', ['year', 'month', 'day', 'hour', 'minute', 'second', 'micros'])
+
+# UTC Datetime Subclasses
+# -----------------------
+def utcnow():
+    return datetime.now(pytz.utc)
+
+class utcdatetime(datetime):
+    def __new__(cls, *args, **kwargs):
+        kwargs['tzinfo'] = pytz.utc
+        dt = datetime.__new__(cls, *args, **kwargs)
+        return dt
+
+# Datetime Calculations
+# ---------------------
+
+WEEKDAYS = [rrule.MO, rrule.TU, rrule.WE, rrule.TH, rrule.FR]
+
+HOLIDAYS = {
+    'new_years'    : datetime(2008 , 1  , 1 ),
+    'mlk_day'      : datetime(2008 , 1  , 21),
+    'presidents'   : datetime(2008 , 2  , 18),
+    'good_friday'  : datetime(2008 , 3  , 21),
+    'memorial_day' : datetime(2008 , 5  , 26),
+    'july_4th'     : datetime(2008 , 7  , 4 ),
+    'labor_day'    : datetime(2008 , 9  , 1 ),
+    'tgiving'      : datetime(2008 , 11 , 27),
+    'christmas'    : datetime(2008 , 5  , 25),
+}
+
+# Create a rule to recur every weekday starting today
+rule = rrule.rrule(
+    rrule.DAILY,
+    byweekday=WEEKDAYS,
+    cache=True,
+)
+
+# Precompute the rule, so that dates are cached.
+rs = rrule.rruleset()
+rs.rrule(rule)
+
+# Add holidays as exclusion days
+for holiday in HOLIDAYS.itervalues():
+    rs.exdate(holiday)
+
+def trading_days(after, before, inclusive=False):
+    """
+    Iterates over the NYSE trading days between the two given
+    dates.
+    """
+    return rs.between(after, before, inc=inclusive)
+
+if __name__ == '__main__':
+
+    now = datetime.now()
+    now30 = datetime.now() + timedelta(days=30)
+
+    # Iterate over the trading days between any two arbitrary
+    # days, excluding the preset holidays.
+    for day in trading_days(now, now30):
+        print day
+
+    # Its now cached so if we do that traversal again it only
+    # takes like 1e-5 seconds.
+    tic = time.time()
+    for day in trading_days(now, now30):
+        print day
+    print time.time() - tic

zipline/finance/performance.py

@@ -194,11 +194,14 @@ class PerformanceTracker():
         Publish the performance results asynchronously to a
         socket.
         """
-        ctx = context or zmq.Context.instance()
-        sock = ctx.socket(zmq.PUSH)
-        sock.connect(zmq_socket)
+        if isinstance(zmq_socket, zmq.Socket):
+            self.result_stream = zmq_socket
+        else:
+            ctx = context or zmq.Context.instance()
+            sock = ctx.socket(zmq.PUSH)
+            sock.connect(zmq_socket)
 
-        self.result_stream = sock
+            self.result_stream = sock
 
     def to_dict(self):
         """

zipline/monitor.py

@@ -169,6 +169,8 @@ class Controller(object):
         self.pub_socket = pub_socket
         self.route_socket = route_socket
 
+        self.error_replay = {}
+
         if logging:
             self.logging = logging
         else:
@@ -242,6 +244,13 @@ class Controller(object):
         #self.logging.info("[Controller] Tracking : %s" % ([c for c in self.tracked],))
         pass
 
+    def replay_errors(self):
+        """
+        Replay the errors in the order they were reported to the
+        controller.
+        """
+        return [ a for a in sorted(self.replay_errors.keys())]
+
     # -------------
     # Publications
     # -------------
@@ -377,6 +386,7 @@ class Controller(object):
         self.logging.info('[Controller] Component "%s" done.' % component)
 
     def exception(self, component, failure):
+        self.error_replay[time.time()] = failure
         self.logging.error('Component "%s" in exception state' % component)
 
     # -----------------

zipline/serial.py

@@ -6,11 +6,9 @@ ZeroMQ. :)
 """
 
 import zlib
-#import blosc
 import hmac
 import base64
-import numpy
-import pandas
+#import blosc
 
 import cPickle as pickle
 
@@ -33,39 +31,6 @@ def recv_zipped_pickle(socket, flags=0, protocol=-1):
     p = zlib.uncompress(z)
     return pickle.loads(p, protocol=protocol)
 
-# HDF5, Numpy Byte Strings, Pandas arrays should use
-# blosc and reconstruct the Python container from the byte string
-# on the other side.
-
-def send_numpy(socket, obj, flags=0):
-    packed = blosc.pack_array(obj)
-    return socket.send(packed, flags=flags)
-
-def recv_numpy(socket, flags=0):
-    packed = blosc.unpack_array(socket.recv(flags))
-    return socket.send(packed, flags=flags)
-
-def send_pandas(socket, obj, flags=0):
-    ndarray = obj._data.blocks[0].values
-    socket.send_multipart(ndarray, flags=flags)
-    spec = (
-        obj._data.index,
-        obj._data.columns,
-        obj._data.blocks[0].dtype
-    )
-    return socket.send_multipart(spec, flags)
-
-def recv_pandas(socket, flags=0):
-    ndarray = socket.recv_multipart(flags)
-    spec = socket.recv_multipart(flags)
-    return pandas.DataFrame._init_ndarray(ndarray, *spec)
-
-def send_hdf5(self):
-    pass
-
-def recv_hdf5(self):
-    pass
-
 # Cryptographically secure wire protocol for ZeroMQ Using HMAC.
 
 # Compare byte strings, backported from Python 3.

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)
+    ])