From 242175350958ed0936aa43461adc4d32d858d4ed Mon Sep 17 00:00:00 2001 From: jfkirk Date: Fri, 26 Jun 2015 18:23:20 -0400 Subject: [PATCH] TST: Fixes broken tests for DataFrameSource --- tests/test_sources.py | 21 +++++---------------- zipline/sources/data_frame_source.py | 8 +++----- 2 files changed, 8 insertions(+), 21 deletions(-) diff --git a/tests/test_sources.py b/tests/test_sources.py index 2e2c7c23..7167a9e3 100644 --- a/tests/test_sources.py +++ b/tests/test_sources.py @@ -14,7 +14,6 @@ # limitations under the License. import pandas as pd import pytz -from itertools import cycle import numpy as np from six import integer_types @@ -26,7 +25,6 @@ from zipline.sources import (DataFrameSource, DataPanelSource, RandomWalkSource) from zipline.utils import tradingcalendar as calendar_nyse -from zipline.finance.trading import with_environment from zipline.assets import AssetFinder @@ -78,25 +76,18 @@ class TestDataFrameSource(TestCase): 'volume', 'price'] copy_panel = data.copy() - copy_panel.items = finder.map_identifier_list_to_sids( + sids = finder.map_identifier_list_to_sids( data.items, data.major_axis[0] ) + copy_panel.items = sids source = DataPanelSource(copy_panel) - sids = [ - asset.sid for asset in - [finder.lookup_symbol(symbol, as_of_date=end) - for symbol in stocks] - ] - stocks_iter = cycle(sids) for event in source: for check_field in check_fields: self.assertIn(check_field, event) self.assertTrue(isinstance(event['volume'], (integer_types))) - self.assertEqual(next(stocks_iter), event['sid']) + self.assertTrue(event['sid'] in sids) - @with_environment() - def test_nan_filter_dataframe(self, env=None): - env.update_asset_finder(identifiers=[4, 5]) + def test_nan_filter_dataframe(self): dates = pd.date_range('1/1/2000', periods=2, freq='B', tz='UTC') df = pd.DataFrame(np.random.randn(2, 2), index=dates, @@ -114,9 +105,7 @@ class TestDataFrameSource(TestCase): self.assertEqual(5, event.sid) self.assertFalse(np.isnan(event.price)) - @with_environment() - def test_nan_filter_panel(self, env=None): - env.update_asset_finder(identifiers=[4, 5]) + def test_nan_filter_panel(self): dates = pd.date_range('1/1/2000', periods=2, freq='B', tz='UTC') df = pd.Panel(np.random.randn(2, 2, 2), major_axis=dates, diff --git a/zipline/sources/data_frame_source.py b/zipline/sources/data_frame_source.py index 77241812..50c0b9b7 100644 --- a/zipline/sources/data_frame_source.py +++ b/zipline/sources/data_frame_source.py @@ -22,7 +22,6 @@ import pandas as pd from zipline.gens.utils import hash_args from zipline.sources.data_source import DataSource -from zipline.finance.trading import with_environment class DataFrameSource(DataSource): @@ -37,9 +36,9 @@ class DataFrameSource(DataSource): Bars where the price is nan are filtered out. """ - @with_environment() - def __init__(self, data, env=None, **kwargs): + def __init__(self, data, **kwargs): assert isinstance(data.index, pd.tseries.index.DatetimeIndex) + # Only accept integer SIDs as the items of the DataFrame assert isinstance(data.columns, pd.Int64Index) # TODO is ffilling correct/necessary? # Forward fill prices @@ -108,8 +107,7 @@ class DataPanelSource(DataSource): Bars where the price is nan are filtered out. """ - @with_environment() - def __init__(self, data, env=None, **kwargs): + def __init__(self, data, **kwargs): assert isinstance(data.major_axis, pd.tseries.index.DatetimeIndex) # Only accept integer SIDs as the items of the Panel assert isinstance(data.items, pd.Int64Index)