From 77fb100ae601a26601f160ea281be258a85cb35e Mon Sep 17 00:00:00 2001 From: warren-oneill Date: Thu, 4 Jun 2015 15:50:41 +0200 Subject: [PATCH] adding type as variable to create_test_panel --- tests/test_sources.py | 4 +++- zipline/finance/performance/tracker.py | 3 ++- zipline/utils/factory.py | 7 ++++++- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/test_sources.py b/tests/test_sources.py index a5e1e833..89ce89a7 100644 --- a/tests/test_sources.py +++ b/tests/test_sources.py @@ -47,14 +47,16 @@ class TestDataFrameSource(TestCase): "DataFrameSource should only stream selected sid 0, not sid 1." def test_panel_source(self): - source, panel = factory.create_test_panel_source() + source, panel = factory.create_test_panel_source(source_type=5) assert isinstance(source.start, pd.lib.Timestamp) assert isinstance(source.end, pd.lib.Timestamp) for event in source: self.assertTrue('sid' in event) self.assertTrue('arbitrary' in event) + self.assertTrue('type' in event) self.assertTrue(hasattr(event, 'volume')) self.assertTrue(hasattr(event, 'price')) + self.assertEquals(event['type'], 5) self.assertEquals(event['arbitrary'], 1.) self.assertEquals(event['sid'], 0) self.assertTrue(isinstance(event['volume'], int)) diff --git a/zipline/finance/performance/tracker.py b/zipline/finance/performance/tracker.py index 96410d7f..6ddf03e6 100644 --- a/zipline/finance/performance/tracker.py +++ b/zipline/finance/performance/tracker.py @@ -339,7 +339,8 @@ class PerformanceTracker(object): def process_close_position(self, event): txn = self.position_tracker.create_close_position_transaction(event) - self.process_transaction(txn) + if txn: + self.process_transaction(txn) def check_upcoming_dividends(self, midnight_of_date_that_just_ended): """ diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index 03da1231..eee575f2 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -315,7 +315,7 @@ def create_test_df_source(sim_params=None, bars='daily'): return DataFrameSource(df), df -def create_test_panel_source(sim_params=None): +def create_test_panel_source(sim_params=None, source_type=None): start = sim_params.first_open \ if sim_params else pd.datetime(1990, 1, 3, 0, 0, 0, 0, pytz.utc) @@ -329,12 +329,17 @@ def create_test_panel_source(sim_params=None): price = np.arange(0, len(index)) volume = np.ones(len(index)) * 1000 + arbitrary = np.ones(len(index)) df = pd.DataFrame({'price': price, 'volume': volume, 'arbitrary': arbitrary}, index=index) + if source_type: + source_types = np.full(len(index), source_type) + df['type'] = source_types + panel = pd.Panel.from_dict({0: df}) return DataPanelSource(panel), panel