diff --git a/tests/test_panel_daily_bar_reader.py b/tests/test_panel_daily_bar_reader.py index 178a4553..e371673a 100644 --- a/tests/test_panel_daily_bar_reader.py +++ b/tests/test_panel_daily_bar_reader.py @@ -13,16 +13,60 @@ # See the License for the specific language governing permissions and # limitations under the License. -from itertools import permutations +from itertools import permutations, product +import numpy as np import pandas as pd from zipline.data.us_equity_pricing import PanelDailyBarReader from zipline.testing import ExplodingObject -from zipline.testing.fixtures import ZiplineTestCase +from zipline.testing.fixtures import ( + WithAssetFinder, + WithNYSETradingDays, + ZiplineTestCase, +) -class TestPanelDailyBarReader(ZiplineTestCase): +class TestPanelDailyBarReader(WithAssetFinder, + WithNYSETradingDays, + ZiplineTestCase): + + START_DATE = pd.Timestamp('2006-01-03', tz='utc') + END_DATE = pd.Timestamp('2006-02-01', tz='utc') + + @classmethod + def init_class_fixtures(cls): + super(TestPanelDailyBarReader, cls).init_class_fixtures() + + finder = cls.asset_finder + days = cls.trading_days + + items = finder.retrieve_all(finder.sids) + major_axis = days + minor_axis = ['open', 'high', 'low', 'close', 'volume'] + + shape = tuple(map(len, [items, major_axis, minor_axis])) + raw_data = np.arange(shape[0] * shape[1] * shape[2]).reshape(shape) + + cls.panel = pd.Panel( + raw_data, + items=items, + major_axis=major_axis, + minor_axis=minor_axis, + ) + + cls.reader = PanelDailyBarReader(days, cls.panel) + + def test_spot_price(self): + panel = self.panel + reader = self.reader + + for asset, date, field in product(*panel.axes): + self.assertEqual( + panel.loc[asset, date, field], + reader.spot_price(asset, date, field), + ) + def test_duplicate_values(self): UNIMPORTANT_VALUE = 57 @@ -37,8 +81,9 @@ class TestPanelDailyBarReader(ZiplineTestCase): axis_names = ['items', 'major_axis', 'minor_axis'] for axis_order in permutations((0, 1, 2)): + transposed = panel.transpose(*axis_order) with self.assertRaises(ValueError) as e: - PanelDailyBarReader(unused, panel.transpose(*axis_order)) + PanelDailyBarReader(unused, transposed) expected = ( "Duplicate entries in Panel.{name}: ['a', 'b'].".format( diff --git a/zipline/data/us_equity_pricing.py b/zipline/data/us_equity_pricing.py index 197b74ca..088f9dce 100644 --- a/zipline/data/us_equity_pricing.py +++ b/zipline/data/us_equity_pricing.py @@ -719,7 +719,7 @@ class PanelDailyBarReader(DailyBarReader): def __init__(self, calendar, panel): panel = panel.copy() - if 'volume' not in panel.items: + if 'volume' not in panel.minor_axis: # Fake volume if it does not exist. panel.loc[:, :, 'volume'] = int(1e9) diff --git a/zipline/utils/input_validation.py b/zipline/utils/input_validation.py index 840ec5b6..96387b87 100644 --- a/zipline/utils/input_validation.py +++ b/zipline/utils/input_validation.py @@ -37,7 +37,8 @@ def verify_indices_all_unique(obj): Returns ------- - None + obj : pd.Series / pd.DataFrame / pd.Panel + The validated object, unchanged. Raises ------ @@ -61,6 +62,7 @@ def verify_indices_all_unique(obj): dupes=sorted(index[index.duplicated()]), ) ) + return obj def optionally(preprocessor):