diff --git a/zipline/algorithm.py b/zipline/algorithm.py index f4f26c02..fca53d4e 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -81,9 +81,8 @@ class TradingAlgorithm(object): def _create_generator(self, environment): """ - Create trading environment, transforms and SimulatedTrading object. - - Gets called by self.run(). + Create a basic generator setup using the sources and + transforms attached to this algorithm. """ self.date_sorted = date_sorted_sources(*self.sources) @@ -96,6 +95,17 @@ class TradingAlgorithm(object): return self.trading_client.simulate(self.with_tnfms) + def get_generator(self, environment): + """ + Override this method to add new logic to the construction + of the generator. Overrides can use the _create_generator + method to get a standard construction generator. + """ + return self._create_generator(environment) + + # 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): """Run the algorithm. diff --git a/zipline/finance/performance.py b/zipline/finance/performance.py index 104caeb0..704383e5 100644 --- a/zipline/finance/performance.py +++ b/zipline/finance/performance.py @@ -404,9 +404,12 @@ class PerformancePeriod(object): self.ending_value = 0.0 self.period_capital_used = 0.0 self.pnl = 0.0 - assert isinstance(initial_positions, positiondict) #sid => position object - self.positions = initial_positions + if not isinstance(initial_positions, positiondict): + self.positions = positiondict() + self.positions.update(initial_positions) + else: + self.positions = initial_positions self.starting_value = starting_value #cash balance at start of period self.starting_cash = starting_cash @@ -435,8 +438,6 @@ class PerformancePeriod(object): def execute_transaction(self, txn): # Update Position # ---------------- - #if txn.sid not in self.positions: - # self.positions[txn.sid] = Position(txn.sid) self.positions[txn.sid].update(txn) self.period_capital_used += -1 * txn.price * txn.amount diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index b94f3c86..0eb6bf20 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -147,10 +147,6 @@ class AlgorithmSimulator(object): # We want an ndict that will have empty ndicts as default # values on missing keys. self.universe = ndict(default=ndict) - # TODO: these keys are being inserted because universe - # has a default dictionary backing __internal. - # del self.universe['__members__'] - # del self.universe['__methods__'] # We don't have a datetime for the current snapshot until we # receive a message. diff --git a/zipline/utils/protocol_utils.py b/zipline/utils/protocol_utils.py index a37895fe..22f4b19a 100644 --- a/zipline/utils/protocol_utils.py +++ b/zipline/utils/protocol_utils.py @@ -47,9 +47,10 @@ class ndict(MutableMapping): def __init__(self, dct=None, default=None): if default is not None: - self.__internal = dict() - else: self.__internal = defaultdict(default) + else: + self.__internal = dict() + if not ndict.cls: ndict.cls = frozenset(dir(self)) diff --git a/zipline/utils/simfactory.py b/zipline/utils/simfactory.py index b9dd78c4..025628eb 100644 --- a/zipline/utils/simfactory.py +++ b/zipline/utils/simfactory.py @@ -86,10 +86,11 @@ def create_test_zipline(**config): concurrent=concurrent_trades ) + test_algo.set_sources([trade_source]) + #------------------- # Transforms #------------------- - test_algo.set_sources([trade_source]) transforms = config.get('transforms', None) if transforms is not None: @@ -104,6 +105,6 @@ def create_test_zipline(**config): # ------------------ # generator/simulator - sim = test_algo._create_generator(trading_environment) + sim = test_algo.get_generator(trading_environment) return sim