""" Generator versions of transforms. """ import random import pytz import logbook import pymongo import types from pymongo import ASCENDING from datetime import datetime, timedelta from collections import deque, defaultdict from numbers import Number from itertools import izip from zipline import ndict from zipline.gens.utils import hash_args, date_gen from zipline.gens.utils import assert_feed_unframe_protocol, assert_transform_protocol import zipline.protocol as zp def PassthroughTransformGen(stream_in): """Trivial transform for event forwarding.""" # hash_args with no arguments is the same as: # hasher = hashlib.md5() # hasher.update(":"); # hashlib.md5.digest(). namestring = "Passthrough" + hash_args() for message in stream_in: assert_feed_unframe_protocol(message) out_value = message assert_transform_protocol(out_value) yield (namestring, out_value) def FunctionalTransformGen(stream_in, fun, *args, **kwargs): """ Generic transform generator that takes each message from an in-stream and yields the output of a function on that message. Not sure how useful this will be in reality, but good for testing. """ # TODO: Distinguish between functions and classes in hash_args. namestring = fun.__name__ + hash_args(*args, **kwargs) for message in stream_in: assert_feed_unframe_protocol(message) out_value = fun(message, *args, **kwargs) assert_transform_protocol(out_value) yield(namestring, out_value) def StatefulTransformGen(stream_in, tnfm_class, *args, **kwargs): """ Generic transform generator that takes each message from an in-stream and feeds it to a state class. For each call to update, the state class must produce a message to be fed downstream. """ # Create an instance of our transform class. state = tnfm_class(*args, **kwargs) # Generate the string associated with this generator's output. namestring = tnfm_class.__name__ + hash_args(*args, **kwargs) for message in stream_in: assert_feed_unframe_protocol(message) out_value = state.update(message) assert_transform_protocol(out_value) yield (namestring, out_value) def MovingAverageTransformGen(stream_in, days, fields): """ Generator that uses the MovingAverage state class to calculate a moving average for all stocks over a specified number of days. """ return StatefulTransformGen(stream_in, MovingAverage, timedelta(days=days), fields) class MovingAverage(object): """ Class that maintains a dictionary from sids to EventWindows Upon receipt of each message we update the corresponding window and return the calculated average. """ def __init__(self, delta, fields): self.delta = delta self.fields = fields # No way to pass arguments to the defaultdict factory, so we # need to define a method to generate the correct EventWindows. self.sid_windows = defaultdict(self.create_window) def create_window(self): """Factory method for self.sid_windows.""" return EventWindow(self.delta, self.fields) def update(self, event): """ Update the event window for this event's sid. Return an ndict from tracked fields to averages. """ assert isinstance(event, ndict),"Bad event in MovingAverage: %s" % event assert event.has_key('sid'), "No sid in MovingAverage: %s" % event # This will create a new EventWindow if this is the first # message for this sid. window = self.sid_windows[event.sid] window.update(event) return window.get_averages() class EventWindow(object): """ Maintains a list of events that are within a certain timedelta of the most recent tick. The expected use of this class is to track events associated with a single sid. We provide simple functionality for averages, but anything more complicated should be handled by a containing class. """ def __init__(self, delta, fields): self.ticks = deque() self.delta = delta self.fields = fields self.totals = defaultdict(float) def __len__(self): return len(self.ticks) def update(self, event): self.assert_well_formed(event) # Add new event and increment totals. self.ticks.append(event) for field in self.fields: self.totals[field] += event[field] # We return a list of all out-of-range events we removed. out_of_range = [] # Clear out expired events, decrementing totals. # newest oldest # | | # V V while (self.ticks[-1].dt - self.ticks[0].dt) >= self.delta: # popleft removes and returns ticks[0] popped = self.ticks.popleft() # Decrement totals for field in self.fields: self.totals[field] -= popped[field] # Add the popped element to the list of dropped events. out_of_range.append(popped) return out_of_range def average(self, field): assert field in self.fields if len(self.ticks) == 0: return 0.0 else: return self.totals[field] / len(self.ticks) def get_averages(self): """ Return an ndict of all our tracked averages. """ out = ndict() for field in self.fields: out[field] = self.average(field) return out def assert_well_formed(self, event): assert isinstance(event, ndict), "Bad event in EventWindow:%s" % event assert event.has_key('dt'), "Missing dt in EventWindow:%s" % event assert isinstance(event.dt, datetime),"Bad dt in EventWindow:%s" % event if len(self.ticks) > 0: # Something is wrong if new event is older than previous. assert event.dt >= self.ticks[-1].dt, \ "Events arrived out of order in EventWindow: %s -> %s" % (event, self.ticks[0]) for field in self.fields: assert event.has_key(field), \ "Event missing [%s] in EventWindow" % field assert isinstance(event[field], Number), \ "Got %s for %s in EventWindow" % (event[field], field) if __name__ == "__main__": def make_event(**kwargs): e = ndict() for key, value in kwargs.iteritems(): e[key] = value return e dates = date_gen(delta = timedelta(hours = 12)) events = ( make_event( sid = 'foo', price = random.random(), dt = date, type = zp.DATASOURCE_TYPE.TRADE, source_id = 'ds', vol = i ) for date, i in izip(dates, xrange(100)) ) gen = MovingAverageTransformGen(events, 1, ['price', 'vol'])