diff --git a/zipline/MESSAGES.py b/zipline/MESSAGES.py index 6d7bd2da..2c67395b 100644 --- a/zipline/MESSAGES.py +++ b/zipline/MESSAGES.py @@ -37,10 +37,4 @@ Please use PerShare or PerTrade. You attempted to override commission after the simulation has \ started. You may only call override_commission in your initialize \ method. -""".strip() - - CALL_RECORD_VARIABLES_POST_INIT = """ -You attempted to register recorded variables after the simulation has \ -started. You may only call record_variables in your initialize \ -method. """.strip() diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 4dd064a6..6ebcc3d4 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -87,7 +87,7 @@ class TradingAlgorithm(object): self.transforms = [] self.sources = [] - self._registered_vars = set() + self._recorded_vars = {} self.logger = None @@ -267,43 +267,16 @@ class TradingAlgorithm(object): 'args': args, 'kwargs': kwargs} - def record_variables(self, names): - """Track and record local variables (i.e. attributes) each - day. - - :Arguments: - names : str or list - List of variable names (strings) to record. - - :Notes: - You are responsible for making sure the attributes - exist. - - The corresponding variable name and its values will be - appended to the results returned by the .run() method. - - :Example: - - In initialize you would call - self.record_variables('mavg'). In handle_data you could - then set self.mavg to some value and it will be recorded. - + def record(self, **kwargs): """ - if not isinstance(names, list): - names = [names] - - for name in names: - if not isinstance(name, basestring): - raise TypeError("record_variables expects only strings") - - if self.initialized: - raise Exception(MESSAGES.ERRORS.CALL_RECORD_VARIABLES_POST_INIT) - - self._registered_vars.update(set(names)) + Track and record local variable (i.e. attributes) each day. + """ + for name, value in kwargs.items(): + self._recorded_vars[name] = value @property def recorded_vars(self): - return {name: getattr(self, name) for name in self._registered_vars} + return copy(self._recorded_vars) @property def portfolio(self): diff --git a/zipline/examples/dual_moving_average.py b/zipline/examples/dual_moving_average.py index fd39eead..8b919ae6 100755 --- a/zipline/examples/dual_moving_average.py +++ b/zipline/examples/dual_moving_average.py @@ -42,8 +42,6 @@ class DualMovingAverage(TradingAlgorithm): # To keep track of whether we invested in the stock or not self.invested = False - self.record_variables(['short_mavg', 'long_mavg', 'buy', 'sell']) - def handle_data(self, data): self.short_mavg = data['AAPL'].short_mavg['price'] self.long_mavg = data['AAPL'].long_mavg['price'] @@ -59,6 +57,11 @@ class DualMovingAverage(TradingAlgorithm): self.invested = False self.sell = True + self.record(short_mavg=self.short_mavg, + long_mavg=self.long_mavg, + buy=self.buy, + sell=self.sell) + if __name__ == '__main__': data = load_from_yahoo(stocks=['AAPL'], indexes={}) dma = DualMovingAverage() diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index 9dd896ab..81de1bed 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -219,10 +219,10 @@ class TimeoutAlgorithm(TradingAlgorithm): class RecordAlgorithm(TradingAlgorithm): def initialize(self): self.incr = 0 - self.record_variables(['incr']) def handle_data(self, data): self.incr += 1 + self.record(incr=self.incr) from zipline.algorithm import TradingAlgorithm from zipline.transforms import BatchTransform, batch_transform