diff --git a/tests/test_components.py b/tests/test_components.py index 45f632fd..ff353186 100644 --- a/tests/test_components.py +++ b/tests/test_components.py @@ -145,7 +145,9 @@ class ComponentTestCase(TestCase): ) launch_monitor(monitor) - sorted_out = date_sorted_sources(comp_a, comp_b, comp_c) + sources = [comp_a, comp_b, comp_c] + gens = [iter(source) for source in sources] + sorted_out = date_sorted_sources(gens) prev = None sort_count = 0 diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py index b4ecde8d..937a9994 100644 --- a/zipline/gens/composites.py +++ b/zipline/gens/composites.py @@ -6,7 +6,7 @@ from zipline.gens.tradegens import SpecificEquityTrades from zipline.gens.utils import roundrobin, hash_args from zipline.gens.sort import date_sort from zipline.gens.merge import merge -from zipline.gens.transform import stateful_transform +from zipline.gens.transform import StatefulTransform SourceBundle = namedtuple("SourceBundle", ['source', 'args', 'kwargs']) TransformBundle = namedtuple("TransformBundle", ['tnfm', 'args', 'kwargs']) @@ -58,8 +58,8 @@ def merged_transforms(sorted_stream, bundles): tnfms_with_streams = zip(split, bundles) # Convert the copies into transform streams. - tnfm_gens = [ - stateful_transform( + tnfms = [ + StatefulTransform( stream_copy, bundle.tnfm, *bundle.args, @@ -67,6 +67,8 @@ def merged_transforms(sorted_stream, bundles): ) for stream_copy, bundle in tnfms_with_streams ] + tnfm_gens = [tnfm.gen() for tnfm in tnfms] + # Roundrobin the outputs of our transforms to create a single flat stream. to_merge = roundrobin(tnfm_gens, namestrings) diff --git a/zipline/gens/examples.py b/zipline/gens/examples.py index 5a24493f..2f7830c8 100644 --- a/zipline/gens/examples.py +++ b/zipline/gens/examples.py @@ -7,7 +7,7 @@ from zipline.test_algorithms import TestAlgorithm from zipline.gens.composites import SourceBundle, TransformBundle, \ date_sorted_sources, merged_transforms from zipline.gens.tradegens import SpecificEquityTrades -from zipline.gens.transform import MovingAverage, Passthrough +from zipline.gens.transform import MovingAverage, Passthrough, StatefulTransform from zipline.gens.tradesimulation import trade_simulation_client as tsc import zipline.protocol as zp @@ -39,11 +39,10 @@ if __name__ == "__main__": sort_out = date_sorted_sources(source_a, source_b) -# passthrough = TransformBundle(Passthrough, (), {}) -# mavg_price = TransformBundle(MovingAverage, (timedelta(minutes = 20), ['price']), {}) -# tnfm_bundles = (passthrough, mavg_price) - -# merge_out = merged_transforms(sort_out, tnfm_bundles) + passthrough = TransformBundle(Passthrough, (), {}) + mavg_price = TransformBundle(MovingAverage, (timedelta(minutes = 20), ['price']), {}) + tnfm_bundles = (passthrough, mavg_price) + merge_out = merged_transforms(sort_out, tnfm_bundles) # # for message in merge_out: # # print message diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 04567fac..6174ac35 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -5,7 +5,7 @@ from numbers import Integral from zipline import ndict -from zipline.gens.transform import stateful_transform +from zipline.gens.transform import StatefulTransform from zipline.finance.trading import TransactionSimulator from zipline.finance.performance import PerformanceTracker diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py index 44ec4b5a..4d453927 100644 --- a/zipline/gens/transform.py +++ b/zipline/gens/transform.py @@ -39,61 +39,72 @@ def functional_transform(stream_in, func, *args, **kwargs): assert_transform_protocol(out_value) yield(namestring, out_value) -def stateful_transform(stream_in, tnfm_class, *args, **kwargs): +class StatefulTransform(object): """ - Generic transform generator that takes each message from an in-stream - and passes it to a state class. For each call to update, the state - class must produce a message to be fed downstream. Any transform class - with the FORWARDER class variable set to true will forward all fields - in the original message. Otherwise only dt, tnfm_id, and tnfm_value - are forwarded. + Generic transform generator that takes each message from an + in-stream and passes it to a state class. For each call to + update, the state class must produce a message to be fed + downstream. Any transform class with the FORWARDER class variable + set to true will forward all fields in the original message. + Otherwise only dt, tnfm_id, and tnfm_value are forwarded. """ - forward_all_fields = tnfm_class.__dict__.get('FORWARDER', False) - update_in_place = tnfm_class.__dict__.get('UPDATER', False) - - assert isinstance(tnfm_class, (types.ObjectType, types.ClassType)), \ + def __init__(self, stream_in, tnfm_class, *args, **kwargs): + assert isinstance(tnfm_class, (types.ObjectType, types.ClassType)), \ "Stateful transform requires a class." - assert tnfm_class.__dict__.has_key('update'), \ + assert tnfm_class.__dict__.has_key('update'), \ "Stateful transform requires the class to have an update method" - - # 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) - - # IMPORTANT: Messages may contain pointers that are shared with - # other streams, so we only manipulate copies. - for message in stream_in: - assert_sort_unframe_protocol(message) - message_copy = deepcopy(message) - - # Same shared pointer issue here as above. - tnfm_value = state.update(deepcopy(message_copy)) - - # If we want to keep all original values, plus append tnfm_id - # and tnfm_value. Used for Passthrough. - if forward_all_fields: - out_message = message_copy - out_message.tnfm_id = namestring - out_message.tnfm_value = tnfm_value - yield out_message + self.forward_all = tnfm_class.__dict__.get('FORWARDER', False) + self.update_in_place = tnfm_class.__dict__.get('UPDATER', False) + assert not all([self.forward_all, self.update_in_place]) - # Our expectation is that the transform simply updated the - # message it was passed. Useful for chaining together - # multiple transforms, e.g. TransactionSimulator/PerformanceTracker. - elif update_in_place: - yield tnfm_value + self.stream_in = stream_in - # Otherwise send tnfm_id, tnfm_value, and the message - # date. Useful for transforms being piped to a merge. - else: - out_message = ndict() - out_message.tnfm_id = namestring - out_message.tnfm_value = tnfm_value - out_message.dt = message_copy.dt - yield out_message + # Create an instance of our transform class. + self.state = tnfm_class(*args, **kwargs) + + # Generate the string associated with this generator's output. + self.namestring = tnfm_class.__name__ + hash_args(*args, **kwargs) + + def get_hash(self): + return self.namestring + + def __iter__(self): + return self.gen() + + def gen(self): + # IMPORTANT: Messages may contain pointers that are shared with + # other streams, so we only manipulate copies. + for message in self.stream_in: + + assert_sort_unframe_protocol(message) + message_copy = deepcopy(message) + + # Same shared pointer issue here as above. + tnfm_value = self.state.update(deepcopy(message_copy)) + + # If we want to keep all original values, plus append tnfm_id + # and tnfm_value. Used for Passthrough. + if self.forward_all: + out_message = message_copy + out_message.tnfm_id = self.namestring + out_message.tnfm_value = tnfm_value + yield out_message + + # Our expectation is that the transform simply updated the + # message it was passed. Useful for chaining together + # multiple transforms, e.g. TransactionSimulator/PerformanceTracker. + elif self.update_in_place: + yield tnfm_value + + # Otherwise send tnfm_id, tnfm_value, and the message + # date. Useful for transforms being piped to a merge. + else: + out_message = ndict() + out_message.tnfm_id = self.namestring + out_message.tnfm_value = tnfm_value + out_message.dt = message_copy.dt + yield out_message class MovingAverage(object): """ diff --git a/zipline/gens/utils.py b/zipline/gens/utils.py index d76966f5..4a95198f 100644 --- a/zipline/gens/utils.py +++ b/zipline/gens/utils.py @@ -43,7 +43,8 @@ def roundrobin(sources, namestrings): """ assert len(sources) == len(namestrings) mapping = OrderedDict(zip(namestrings, sources)) - + + import nose.tools; nose.tools.set_trace() # While our generators have not been exhausted, pull elements while mapping.keys() != []: for namestring, source in mapping.iteritems():