diff --git a/tests/test_algorithm_gen.py b/tests/test_algorithm_gen.py new file mode 100644 index 00000000..0d92ce7e --- /dev/null +++ b/tests/test_algorithm_gen.py @@ -0,0 +1,103 @@ +#!/usr/bin/python +# +# Copyright 2012 Quantopian, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import TestCase +from nose.tools import timed + +from datetime import datetime +import pytz + +from zipline.algorithm import TradingAlgorithm +from zipline.finance import slippage +from zipline.utils import factory +from zipline.utils.test_utils import ( + setup_logger, + teardown_logger +) + +DEFAULT_TIMEOUT = 15 # seconds +EXTENDED_TIMEOUT = 90 + + +class RecordDateSlippage(slippage.FixedSlippage): + def __init__(self, spread): + super(RecordDateSlippage, self).__init__(spread=spread) + self.latest_date = None + + def simulate(self, event, open_orders): + self.latest_date = event['datetime'] + result = super(RecordDateSlippage, self).simulate(event, open_orders) + return result + + +class TestAlgo(TradingAlgorithm): + + def __init__(self, asserter, *args, **kwargs): + super(TestAlgo, self).__init__(*args, **kwargs) + self.asserter = asserter + + def initialize(self, window_length=100): + self.latest_date = None + + self.set_slippage(RecordDateSlippage(spread=0.05)) + self.stocks = [8229] + self.ordered = False + + def handle_data(self, data): + self.latest_date = self.get_datetime() + + if not self.ordered: + for stock in self.stocks: + self.order(stock, 100) + + self.ordered = True + + self.asserter.assertGreaterEqual( + self.latest_date, + self.slippage.latest_date + ) + + +class AlgorithmGeneratorTestCase(TestCase): + def setUp(self): + setup_logger(self) + + def tearDown(self): + teardown_logger(self) + + @timed(DEFAULT_TIMEOUT) + def test_generator_dates(self): + """ + Ensure the pipeline of generators are in sync, at least as far as + their current dates. + """ + algo = TestAlgo(self) + trading_environment = factory.create_trading_environment( + start=datetime(2012, 1, 3, tzinfo=pytz.utc), + end=datetime(2012, 7, 30, tzinfo=pytz.utc) + ) + trade_source = factory.create_daily_trade_source( + [8229], + 200, + trading_environment + ) + algo.set_sources([trade_source]) + + gen = algo.get_generator(trading_environment) + self.assertTrue(list(gen)) + + self.assertTrue(algo.slippage.latest_date) + self.assertTrue(algo.latest_date) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 463851a9..c55fc38d 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -20,6 +20,9 @@ import numpy as np from datetime import datetime +from itertools import groupby +from operator import attrgetter + from zipline.sources import DataFrameSource from zipline.utils.factory import create_trading_environment from zipline.transforms.utils import StatefulTransform @@ -95,12 +98,15 @@ class TradingAlgorithm(object): self.date_sorted = date_sorted_sources(*self.sources) self.with_tnfms = sequential_transforms(self.date_sorted, *self.transforms) + # Group together events with the same dt field. This depends on the + # events already being sorted. + self.grouped_by_date = groupby(self.with_tnfms, attrgetter('dt')) self.trading_client = tsc(self, environment) transact_method = transact_partial(self.slippage, self.commission) self.set_transact(transact_method) - return self.trading_client.simulate(self.with_tnfms) + return self.trading_client.simulate(self.grouped_by_date) def get_generator(self, environment): """ diff --git a/zipline/finance/performance.py b/zipline/finance/performance.py index 98e2ac8d..532d508f 100644 --- a/zipline/finance/performance.py +++ b/zipline/finance/performance.py @@ -210,16 +210,18 @@ class PerformanceTracker(object): """ Main generator work loop. """ - for event in stream_in: - if event.dt == "DONE": - event.perf_message = self.handle_simulation_end() - del event['TRANSACTION'] - yield event - else: - event.perf_message = self.process_event(event) - event.portfolio = self.get_portfolio() - del event['TRANSACTION'] - yield event + for date, snapshot in stream_in: + yield date, [self._transform_event(event) for event in snapshot] + + def _transform_event(self, event): + if event.dt == "DONE": + event.perf_message = self.handle_simulation_end() + else: + event.perf_message = self.process_event(event) + event.portfolio = self.get_portfolio() + + del event['TRANSACTION'] + return event def get_portfolio(self): return self.cumulative_performance.as_portfolio() diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index 502d0637..5d7eaae8 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -44,8 +44,8 @@ class TransactionSimulator(object): """ Main generator work loop. """ - for event in stream_in: - yield self.update(event) + for date, snapshot in stream_in: + yield date, [self.update(event) for event in snapshot] def update(self, event): event.TRANSACTION = None diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index b9fcda85..5a70746c 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -18,8 +18,6 @@ from logbook import Logger, Processor from collections import defaultdict from datetime import datetime -from itertools import groupby -from operator import attrgetter from zipline import ndict @@ -200,9 +198,8 @@ class AlgorithmSimulator(object): # inject the current algo # snapshot time to any log record generated. with self.processor.threadbound(): - # Group together events with the same dt field. This depends on the - # events already being sorted. - for date, snapshot in groupby(stream_in, attrgetter('dt')): + + for date, snapshot in stream_in: # Set the simulation date to be the first event we see. # This should only occur once, at the start of the test. if self.simulation_dt is None: