diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index be953be2..76b3ddc0 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -25,6 +25,7 @@ from zipline.sources import (SpecificEquityTrades, DataFrameSource, DataPanelSource) from zipline.transforms import MovingAverage +from zipline.finance.trading import SimulationParameters class TestRecordAlgorithm(TestCase): @@ -88,12 +89,15 @@ class TestTransformAlgorithm(TestCase): algo.run([self.source, self.df_source]) def test_multi_source_as_input(self): + sim_params = SimulationParameters( + self.df.index[0], + self.df.index[-1] + ) algo = TestRegisterTransformAlgorithm( - self.sim_params, + sim_params=sim_params, sids=[0, 1, 133] ) - algo.run([self.source, self.df_source], - start=self.df.index[0], end=self.df.index[-1]) + algo.run([self.source, self.df_source]) self.assertEqual(len(algo.sources), 2) def test_df_as_input(self): diff --git a/zipline/algorithm.py b/zipline/algorithm.py index b4005946..4dd064a6 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -150,7 +150,7 @@ class TradingAlgorithm(object): # TODO: make a new subclass, e.g. BatchAlgorithm, and move # the run method to the subclass, and refactor to put the # generator creation logic into get_generator. - def run(self, source, start=None, end=None): + def run(self, source, sim_params=None): """Run the algorithm. :Arguments: @@ -172,9 +172,10 @@ class TradingAlgorithm(object): """ if isinstance(source, (list, tuple)): - assert start is not None and end is not None, \ + assert self.sim_params is not None or sim_params is not None, \ """When providing a list of sources, \ - start and end date have to be specified.""" + sim_params have to be specified as a parameter + or in the constructor.""" elif isinstance(source, pd.DataFrame): # if DataFrame provided, wrap in DataFrameSource source = DataFrameSource(source) @@ -182,9 +183,8 @@ class TradingAlgorithm(object): source = DataPanelSource(source) # If values not set, try to extract from source. - if start is None: + if self.sim_params is None and sim_params is None: start = source.start - if end is None: end = source.end if not isinstance(source, (list, tuple)): @@ -192,6 +192,9 @@ class TradingAlgorithm(object): else: self.sources = source + if sim_params: + self.sim_params = sim_params + if not self.sim_params: self.sim_params = create_simulation_parameters( start=start,