From dba86153d23ef925761738b531c8c23ad26a028f Mon Sep 17 00:00:00 2001 From: fawce Date: Fri, 8 Mar 2013 22:59:35 -0500 Subject: [PATCH] ENH: added a CUSTOM datasource type for custom data. - perf modified to let non-performance related events flow through. - changes to support streaming non-trading data through batch transforms and for mixing in sids with just custom data. - allowing CUSTOM events to flow through to transforms. - Added logic to maintain pre-specified sid filter. --- zipline/finance/performance.py | 6 ++++ zipline/gens/tradesimulation.py | 1 - zipline/protocol.py | 3 +- zipline/transforms/utils.py | 63 +++++++++++++++++++++++++-------- 4 files changed, 57 insertions(+), 16 deletions(-) diff --git a/zipline/finance/performance.py b/zipline/finance/performance.py index 2c239146..e299d077 100644 --- a/zipline/finance/performance.py +++ b/zipline/finance/performance.py @@ -273,6 +273,12 @@ class PerformanceTracker(object): elif event.type == zp.DATASOURCE_TYPE.DIVIDEND: self.cumulative_performance.add_dividend(event) self.todays_performance.add_dividend(event) + # this event will not be relayed up + messages = None + elif event.type == zp.DATASOURCE_TYPE.CUSTOM: + # we just want to relay this event unchanged. + messages = [] + return messages #calculate performance as of last trade self.cumulative_performance.calculate_performance() diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 576de754..f3d54352 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -273,7 +273,6 @@ class AlgorithmSimulator(object): # of the event data sent to handle_data. To avoid # confusion, we remove it from the event here. del event.portfolio - # Update our knowledge of this event's sid sid_data = self.universe[event.sid] sid_data.__dict__.update(event.__dict__) diff --git a/zipline/protocol.py b/zipline/protocol.py index c6cf2982..38ba7bf2 100644 --- a/zipline/protocol.py +++ b/zipline/protocol.py @@ -27,7 +27,8 @@ DATASOURCE_TYPE = Enum( 'DIVIDEND', 'TRADE', 'EMPTY', - 'DONE' + 'DONE', + 'CUSTOM' ) diff --git a/zipline/transforms/utils.py b/zipline/transforms/utils.py index 7b47d9f7..2c44dbcd 100644 --- a/zipline/transforms/utils.py +++ b/zipline/transforms/utils.py @@ -137,7 +137,12 @@ class StatefulTransform(object): for message in stream_in: # we only handle TRADE events. if (hasattr(message, 'type') - and message.type != DATASOURCE_TYPE.TRADE): + and message.type not in ( + DATASOURCE_TYPE.TRADE, + DATASOURCE_TYPE.CUSTOM)): + # TODO: this should be yielding the original message + # instead of swallowing it. Will be an issue when we + # have a transaction source from brokers etc. continue # allow upstream generators to yield None to avoid # blocking. @@ -218,7 +223,10 @@ class EventWindow(object): def update(self, event): - if hasattr(event, 'type') and event.type != DATASOURCE_TYPE.TRADE: + if (hasattr(event, 'type') + and event.type not in ( + DATASOURCE_TYPE.TRADE, + DATASOURCE_TYPE.CUSTOM)): return self.assert_well_formed(event) @@ -354,13 +362,25 @@ class BatchTransform(EventWindow): self.clean_nans = clean_nans self.compute_only_full = compute_only_full - self.sids = sids - if isinstance(self.sids, (basestring, Integral)): - self.sids = [self.sids] + # The following logic is to allow pre-specified sid filters + # to operate on the data, but to also allow new symbols to + # enter the batch transform's window IFF a sid filter is not + # specified. + self.sids = None + if sids: + self.static_sids = True + self.sids = sids + if isinstance(sids, (basestring, Integral)): + self.sids = set([sids]) + else: + self.sids = set(sids) + else: + self.static_sids = False - self.field_names = fields - if isinstance(self.field_names, str): - self.field_names = [self.field_names] + self.initial_field_names = fields + if isinstance(self.initial_field_names, basestring): + self.initial_field_names = [self.initial_field_names] + self.field_names = set() self.refresh_period = refresh_period self.window_length = window_length @@ -430,19 +450,34 @@ class BatchTransform(EventWindow): if (isinstance(value, (int, float)))]) sid_keys.append(keys) - assert sid_keys[0] == set.intersection(*sid_keys),\ - "Each sid must have the same keys." - + # with CUSTOM data events, there may be different fields + # per sid. So the allowable keys are the union of all events. + union = set.union(*sid_keys) unwanted_fields = set(['portfolio', 'sid', 'dt', 'type', 'datetime', 'source_id']) - return sid_keys[0] - unwanted_fields + return union - unwanted_fields def handle_add(self, event): if not self.last_dt: - if self.field_names is None: - self.field_names = self._extract_field_names(event) self.last_dt = event.dt + if self.initial_field_names is None: + self.latest_names = self._extract_field_names(event) + if self.field_names: + self.field_names = \ + set.union(self.field_names, self.latest_names) + else: + self.field_names = self.latest_names + else: + self.field_names = self.initial_field_names + + if not self.static_sids: + if self.sids: + event_sids = set(event.data.keys()) + self.sids = set.union(self.sids, event_sids) + else: + self.sids = set(event.data.keys()) + # update trading day counters if self.last_dt.day != event.dt.day: self.last_dt = event.dt