From c0eb798cc65d74696116b4a98c3abec14b818972 Mon Sep 17 00:00:00 2001 From: Maya Tydykov Date: Tue, 10 May 2016 12:46:53 -0400 Subject: [PATCH] TST: modify test class to use WithAssetFinder fixture. BUG: assign result to var TST: remove obsolete assertion STY: fix line length --- tests/pipeline/test_blaze.py | 441 ++++++++++++++----------- zipline/pipeline/loaders/blaze/core.py | 76 ++++- 2 files changed, 304 insertions(+), 213 deletions(-) diff --git a/tests/pipeline/test_blaze.py b/tests/pipeline/test_blaze.py index a0faf6be..512be3ee 100644 --- a/tests/pipeline/test_blaze.py +++ b/tests/pipeline/test_blaze.py @@ -6,7 +6,6 @@ from __future__ import division from collections import OrderedDict from datetime import timedelta, time from itertools import product, chain -from unittest import TestCase import warnings import blaze as bz @@ -33,12 +32,13 @@ from zipline.pipeline.loaders.blaze.core import ( NonPipelineField, no_deltas_rules, ) +from zipline.testing.fixtures import WithAssetFinder from zipline.utils.numpy_utils import ( float64_dtype, int64_dtype, repeat_last_axis, ) -from zipline.testing import tmp_asset_finder +from zipline.testing import tmp_asset_finder, ZiplineTestCase nameof = op.attrgetter('name') dtypeof = op.attrgetter('dtype') @@ -73,14 +73,17 @@ def _utc_localize_index_level_0(df): return df -class BlazeToPipelineTestCase(TestCase): +class BlazeToPipelineTestCase(WithAssetFinder, ZiplineTestCase): + START_DATE = pd.Timestamp(0) + END_DATE = pd.Timestamp('2015') + @classmethod - def setUpClass(cls): + def init_class_fixtures(cls): + super(BlazeToPipelineTestCase, cls).init_class_fixtures() cls.dates = dates = pd.date_range('2014-01-01', '2014-01-03') dates = cls.dates.repeat(3) - cls.sids = sids = ord('A'), ord('B'), ord('C') cls.df = df = pd.DataFrame({ - 'sid': sids * 3, + 'sid': cls.ASSET_FINDER_EQUITY_SIDS * 3, 'value': (0., 1., 2., 1., 2., 3., 2., 3., 4.), 'int_value': (0, 1, 2, 1, 2, 3, 2, 3, 4), 'asof_date': dates, @@ -121,14 +124,6 @@ class BlazeToPipelineTestCase(TestCase): self.assertTrue(np.isnan(ds.value.missing_value)) self.assertEqual(ds.int_value.missing_value, 0) - invalid_type_fields = ('asof_date',) - - for field in invalid_type_fields: - with self.assertRaises(AttributeError) as e: - getattr(ds, field) - self.assertIn("'%s'" % field, str(e.exception)) - self.assertIn("'datetime64[us]'", str(e.exception)) - # test memoization self.assertIs( from_blaze( @@ -276,11 +271,15 @@ class BlazeToPipelineTestCase(TestCase): ) def test_cols_with_missing_vals(self): - dates = (pd.Timestamp('2014-01-01'), pd.Timestamp('2014-01-03')) + dates = (self.dates[0], self.dates[-1]) df = pd.DataFrame({ - 'sid': self.sids[:-1], - 'value': (0., 1.,), + 'sid': self.ASSET_FINDER_EQUITY_SIDS[:-1], + 'float_value': (0., 1.,), 'str_value': ("a", "b",), + 'int_value': (1, 2), + 'bool_value': (True, True), + 'dt_value': (pd.Timestamp('2011-01-01'), + pd.Timestamp('2011-01-02')), 'asof_date': dates, 'timestamp': dates, }) @@ -289,37 +288,85 @@ class BlazeToPipelineTestCase(TestCase): dshape=""" var * { sid: int64, - value: float64, + float_value: float64, str_value: string, + int_value: int64, + bool_value: bool, + dt_value: datetime, asof_date: datetime, timestamp: datetime, }""", ) fields = OrderedDict(expr.dshape.measure.fields) - with tmp_asset_finder() as finder: - expected = pd.DataFrame( - np.array([["a", 0], - [None, np.NaN], - [None, np.NaN], - ["a", 0], - [None, np.NaN], - [None, np.NaN], - ["a", 0], - ["b", 1], - [None, np.NaN]]), - columns=['str_value', 'value'], - index=pd.MultiIndex.from_product( - (self.dates, finder.retrieve_all(self.sids)) - ) - ) - self._test_id( - df, - var * Record(fields), - expected, - finder, - ('value', 'str_value'), + expected = pd.DataFrame({ + "str_value": np.array(["a", + None, + None, + "a", + None, + None, + "a", + "b", + None], + dtype='object'), + "float_value": np.array([0, + np.NaN, + np.NaN, + 0, + np.NaN, + np.NaN, + 0, + 1, + np.NaN], + dtype='float64'), + "int_value": np.array([1, + 0, + 0, + 1, + 0, + 0, + 1, + 2, + 0], + dtype='int64'), + "bool_value": np.array([True, + False, + False, + True, + False, + False, + True, + True, + False], + dtype='bool'), + "dt_value": [pd.Timestamp('2011-01-01'), + pd.NaT, + pd.NaT, + pd.Timestamp('2011-01-01'), + pd.NaT, + pd.NaT, + pd.Timestamp('2011-01-01'), + pd.Timestamp('2011-01-02'), + pd.NaT], + }, + columns=['str_value', 'float_value', 'int_value', 'bool_value', + 'dt_value'], + index=pd.MultiIndex.from_product( + (self.dates, self.asset_finder.retrieve_all( + self.ASSET_FINDER_EQUITY_SIDS + )) ) + ) + + self._test_id( + df, + var * Record(fields), + expected, + self.asset_finder, + ('float_value', 'str_value', 'int_value', 'bool_value', + 'dt_value'), + ) def test_complex_expr(self): expr = bz.data(self.df, dshape=self.dshape) @@ -376,15 +423,14 @@ class BlazeToPipelineTestCase(TestCase): p.add(getattr(ds, a).latest, a) dates = self.dates - with tmp_asset_finder() as finder: - result = SimplePipelineEngine( - loader, - dates, - finder, - ).run_pipeline(p, dates[0], dates[-1]) + result = SimplePipelineEngine( + loader, + dates, + finder, + ).run_pipeline(p, dates[0], dates[-1]) assert_frame_equal( - result, - _utc_localize_index_level_0(expected), + result.sort_index(axis=1), + _utc_localize_index_level_0(expected.sort_index(axis=1)), check_dtype=False, ) @@ -408,12 +454,11 @@ class BlazeToPipelineTestCase(TestCase): p.add(ds.int_value.latest, 'int_value') dates = self.dates - with tmp_asset_finder() as finder: - result = SimplePipelineEngine( - loader, - dates, - finder, - ).run_pipeline(p, dates[0], dates[-1]) + result = SimplePipelineEngine( + loader, + dates, + self.asset_finder, + ).run_pipeline(p, dates[0], dates[-1]) expected = df.drop('asof_date', axis=1) expected['timestamp'] = expected['timestamp'].dt.normalize().astype( @@ -423,7 +468,7 @@ class BlazeToPipelineTestCase(TestCase): expected.set_index(['timestamp', 'sid'], inplace=True) expected.index = pd.MultiIndex.from_product(( expected.index.levels[0], - finder.retrieve_all(expected.index.levels[1]), + self.asset_finder.retrieve_all(expected.index.levels[1]), )) assert_frame_equal(result, expected, check_dtype=False) @@ -453,17 +498,17 @@ class BlazeToPipelineTestCase(TestCase): Equity(66 [B]) 3 Equity(67 [C]) 4 """ - with tmp_asset_finder() as finder: - expected = self.df.drop('asof_date', axis=1).set_index( - ['timestamp', 'sid'], - ) - expected.index = pd.MultiIndex.from_product(( - expected.index.levels[0], - finder.retrieve_all(expected.index.levels[1]), - )) - self._test_id( - self.df, self.dshape, expected, finder, ('int_value', 'value',) - ) + expected = self.df.drop('asof_date', axis=1).set_index( + ['timestamp', 'sid'], + ) + expected.index = pd.MultiIndex.from_product(( + expected.index.levels[0], + self.asset_finder.retrieve_all(expected.index.levels[1]), + )) + self._test_id( + self.df, self.dshape, expected, self.asset_finder, + ('int_value', 'value',) + ) def test_id_ffill_out_of_window(self): """ @@ -494,7 +539,7 @@ class BlazeToPipelineTestCase(TestCase): """ dates = self.dates.repeat(3) - timedelta(days=10) df = pd.DataFrame({ - 'sid': self.sids * 3, + 'sid': self.ASSET_FINDER_EQUITY_SIDS * 3, 'value': (0, 1, np.nan, 1, np.nan, 3, np.nan, 3, 4), 'other': (0, np.nan, 2, np.nan, 2, 3, 2, 3, np.nan), 'asof_date': dates, @@ -503,29 +548,30 @@ class BlazeToPipelineTestCase(TestCase): fields = OrderedDict(self.dshape.measure.fields) fields['other'] = fields['value'] - with tmp_asset_finder() as finder: - expected = pd.DataFrame( - np.array([[2, 1], - [3, 3], - [3, 4], - [2, 1], - [3, 3], - [3, 4], - [2, 1], - [3, 3], - [3, 4]]), - columns=['other', 'value'], - index=pd.MultiIndex.from_product( - (self.dates, finder.retrieve_all(self.sids)), - ), - ) - self._test_id( - df, - var * Record(fields), - expected, - finder, - ('value', 'other'), - ) + expected = pd.DataFrame( + np.array([[2, 1], + [3, 3], + [3, 4], + [2, 1], + [3, 3], + [3, 4], + [2, 1], + [3, 3], + [3, 4]]), + columns=['other', 'value'], + index=pd.MultiIndex.from_product( + (self.dates, self.asset_finder.retrieve_all( + self.ASSET_FINDER_EQUITY_SIDS + )), + ), + ) + self._test_id( + df, + var * Record(fields), + expected, + self.asset_finder, + ('value', 'other'), + ) def test_id_multiple_columns(self): """ @@ -557,21 +603,20 @@ class BlazeToPipelineTestCase(TestCase): df['other'] = df.value + 1 fields = OrderedDict(self.dshape.measure.fields) fields['other'] = fields['value'] - with tmp_asset_finder() as finder: - expected = df.drop('asof_date', axis=1).set_index( - ['timestamp', 'sid'], - ).sort_index(axis=1) - expected.index = pd.MultiIndex.from_product(( - expected.index.levels[0], - finder.retrieve_all(expected.index.levels[1]), - )) - self._test_id( - df, - var * Record(fields), - expected, - finder, - ('value', 'int_value', 'other'), - ) + expected = df.drop('asof_date', axis=1).set_index( + ['timestamp', 'sid'], + ).sort_index(axis=1) + expected.index = pd.MultiIndex.from_product(( + expected.index.levels[0], + self.asset_finder.retrieve_all(expected.index.levels[1]), + )) + self._test_id( + df, + var * Record(fields), + expected, + self.asset_finder, + ('value', 'int_value', 'other'), + ) def test_id_macro_dataset(self): """ @@ -595,22 +640,21 @@ class BlazeToPipelineTestCase(TestCase): """ asset_info = asset_infos[0][0] nassets = len(asset_info) - with tmp_asset_finder() as finder: - expected = pd.DataFrame( - list(concatv([0] * nassets, [1] * nassets, [2] * nassets)), - index=pd.MultiIndex.from_product(( - self.macro_df.timestamp, - finder.retrieve_all(asset_info.index), - )), - columns=('value',), - ) - self._test_id( - self.macro_df, - self.macro_dshape, - expected, - finder, - ('value',), - ) + expected = pd.DataFrame( + list(concatv([0] * nassets, [1] * nassets, [2] * nassets)), + index=pd.MultiIndex.from_product(( + self.macro_df.timestamp, + self.asset_finder.retrieve_all(asset_info.index), + )), + columns=('value',), + ) + self._test_id( + self.macro_df, + self.macro_dshape, + expected, + self.asset_finder, + ('value',), + ) def test_id_ffill_out_of_window_macro_dataset(self): """ @@ -642,29 +686,30 @@ class BlazeToPipelineTestCase(TestCase): fields = OrderedDict(self.macro_dshape.measure.fields) fields['other'] = fields['value'] - with tmp_asset_finder() as finder: - expected = pd.DataFrame( - np.array([[0, 1], - [0, 1], - [0, 1], - [0, 1], - [0, 1], - [0, 1], - [0, 1], - [0, 1], - [0, 1]]), - columns=['value', 'other'], - index=pd.MultiIndex.from_product( - (self.dates, finder.retrieve_all(self.sids)), - ), - ).sort_index(axis=1) - self._test_id( - df, - var * Record(fields), - expected, - finder, - ('value', 'other'), - ) + expected = pd.DataFrame( + np.array([[0, 1], + [0, 1], + [0, 1], + [0, 1], + [0, 1], + [0, 1], + [0, 1], + [0, 1], + [0, 1]]), + columns=['value', 'other'], + index=pd.MultiIndex.from_product( + (self.dates, self.asset_finder.retrieve_all( + self.ASSET_FINDER_EQUITY_SIDS + )), + ), + ).sort_index(axis=1) + self._test_id( + df, + var * Record(fields), + expected, + self.asset_finder, + ('value', 'other'), + ) def test_id_macro_dataset_multiple_columns(self): """ @@ -739,31 +784,32 @@ class BlazeToPipelineTestCase(TestCase): fields = OrderedDict(self.dshape.measure.fields) fields['other'] = fields['value'] - with tmp_asset_finder() as finder: - expected = pd.DataFrame( - columns=['other', 'value'], - data=[ - [1, 0], # 2014-01-01 Equity(65 [A]) - [np.nan, 1], # Equity(66 [B]) - [2, np.nan], # Equity(67 [C]) - [1, 1], # 2014-01-02 Equity(65 [A]) - [2, 1], # Equity(66 [B]) - [3, 3], # Equity(67 [C]) - [2, 1], # 2014-01-03 Equity(65 [A]) - [3, 3], # Equity(66 [B]) - [3, 3], # Equity(67 [C]) - ], - index=pd.MultiIndex.from_product( - (self.dates, finder.retrieve_all(self.sids)), - ), - ) - self._test_id( - df, - var * Record(fields), - expected, - finder, - ('value', 'other'), - ) + expected = pd.DataFrame( + columns=['other', 'value'], + data=[ + [1, 0], # 2014-01-01 Equity(65 [A]) + [np.nan, 1], # Equity(66 [B]) + [2, np.nan], # Equity(67 [C]) + [1, 1], # 2014-01-02 Equity(65 [A]) + [2, 1], # Equity(66 [B]) + [3, 3], # Equity(67 [C]) + [2, 1], # 2014-01-03 Equity(65 [A]) + [3, 3], # Equity(66 [B]) + [3, 3], # Equity(67 [C]) + ], + index=pd.MultiIndex.from_product( + (self.dates, self.asset_finder.retrieve_all( + self.ASSET_FINDER_EQUITY_SIDS + )), + ), + ) + self._test_id( + df, + var * Record(fields), + expected, + self.asset_finder, + ('value', 'other'), + ) def test_id_take_last_in_group_macro(self): """ @@ -795,33 +841,34 @@ class BlazeToPipelineTestCase(TestCase): fields = OrderedDict(self.macro_dshape.measure.fields) fields['other'] = fields['value'] - with tmp_asset_finder() as finder: - expected = pd.DataFrame( - columns=[ - 'other', 'value', - ], - data=[ - [np.nan, 1], # 2014-01-01 Equity(65 [A]) - [np.nan, 1], # Equity(66 [B]) - [np.nan, 1], # Equity(67 [C]) - [1, 2], # 2014-01-02 Equity(65 [A]) - [1, 2], # Equity(66 [B]) - [1, 2], # Equity(67 [C]) - [2, 2], # 2014-01-03 Equity(65 [A]) - [2, 2], # Equity(66 [B]) - [2, 2], # Equity(67 [C]) - ], - index=pd.MultiIndex.from_product( - (self.dates, finder.retrieve_all(self.sids)), - ), - ) - self._test_id( - df, - var * Record(fields), - expected, - finder, - ('value', 'other'), - ) + expected = pd.DataFrame( + columns=[ + 'other', 'value', + ], + data=[ + [np.nan, 1], # 2014-01-01 Equity(65 [A]) + [np.nan, 1], # Equity(66 [B]) + [np.nan, 1], # Equity(67 [C]) + [1, 2], # 2014-01-02 Equity(65 [A]) + [1, 2], # Equity(66 [B]) + [1, 2], # Equity(67 [C]) + [2, 2], # 2014-01-03 Equity(65 [A]) + [2, 2], # Equity(66 [B]) + [2, 2], # Equity(67 [C]) + ], + index=pd.MultiIndex.from_product( + (self.dates, self.asset_finder.retrieve_all( + self.ASSET_FINDER_EQUITY_SIDS + )), + ), + ) + self._test_id( + df, + var * Record(fields), + expected, + self.asset_finder, + ('value', 'other'), + ) def _run_pipeline(self, expr, @@ -1037,7 +1084,7 @@ class BlazeToPipelineTestCase(TestCase): ]) repeated_dates = base_dates.repeat(3) baseline = pd.DataFrame({ - 'sid': self.sids * 2, + 'sid': self.ASSET_FINDER_EQUITY_SIDS * 2, 'value': (0., 1., 2., 1., 2., 3.), 'int_value': (0, 1, 2, 1, 2, 3), 'asof_date': repeated_dates, diff --git a/zipline/pipeline/loaders/blaze/core.py b/zipline/pipeline/loaders/blaze/core.py index 53504fe7..f676bcf7 100644 --- a/zipline/pipeline/loaders/blaze/core.py +++ b/zipline/pipeline/loaders/blaze/core.py @@ -140,7 +140,6 @@ from datashape import ( floating, isrecord, isscalar, - String ) import numpy as np from odo import odo @@ -177,7 +176,11 @@ from zipline.utils.input_validation import ( ensure_timezone, optionally, ) -from zipline.utils.numpy_utils import categorical_dtype, repeat_last_axis +from zipline.utils.numpy_utils import ( + categorical_dtype, + repeat_last_axis, + datetime64ns_dtype +) from zipline.utils.pandas_utils import sort_values from zipline.utils.preprocess import preprocess @@ -275,6 +278,31 @@ class NotPipelineCompatible(TypeError): _new_names = ('BlazeDataSet_%d' % n for n in count()) +def datashape_type_to_numpy(type_): + """ + Given a datashape type, return the associated numpy type. Maps + datashape's DateTime type to numpy's `datetime64[ns]` dtype, since the + numpy datetime returned by datashape isn't supported by pipeline. + + Parameters + ---------- + type_: datashape.coretypes.Type + The datashape type. + + Returns + ------- + type_ np.dtype + The numpy dtype. + + """ + if isinstance(type_, Option): + type_ = type_.ty + if isinstance(type_, DateTime): + return np.dtype('datetime64[ns]') + else: + return type_.to_numpy_dtype() + + @memoize def new_dataset(expr, deltas, missing_values): """ @@ -310,17 +338,13 @@ def new_dataset(expr, deltas, missing_values): # Terms. if name in (SID_FIELD_NAME, TS_FIELD_NAME): continue - try: - if isinstance(type_, Option): - type_ = type_.ty - type_ = type_.to_numpy_dtype() - if not can_represent_dtype(type_): - raise NotPipelineCompatible() + type_ = datashape_type_to_numpy(type_) + if can_represent_dtype(type_): col = Column( type_, missing_values.get(name, NotSpecified), ) - except NotPipelineCompatible: + else: col = NonPipelineField(name, type_) columns[name] = col @@ -1007,19 +1031,39 @@ class BlazeLoader(dict): else: last_in_group = last_in_group.reindex(dates) # Unstack will fill all missing values with NaN; we need to fix - # this for strings. + # this for all types that are not float. if not df.empty: - str_cols = df.columns[df.dtypes == categorical_dtype] - - for col in str_cols: - last_in_group[col] = last_in_group[col].where(pd.notnull( - last_in_group[col]), None) + for column in columns: + if df[column.name].dtype == categorical_dtype: + last_in_group[column.name] = last_in_group[ + column.name + ].where(pd.notnull(last_in_group[column.name]), + column.missing_value) + # Need to convert from float col to datetime col + elif df[column.name].dtype == datetime64ns_dtype: + last_in_group[column.name] = last_in_group[ + column.name + ].astype('datetime64[ns]') + else: + last_in_group[column.name] = last_in_group[ + column.name + ].fillna(column.missing_value) return last_in_group sparse_deltas = last_in_date_group(non_novel_deltas, reindex=False) dense_output = last_in_date_group(sparse_output, reindex=True) - dense_output.ffill(inplace=True) + for column in columns: + if have_sids: + dense_output[column.name] = dense_output[ + column.name + ].apply(lambda x: x.replace( + to_replace=column.missing_value, method='ffill' + )) + else: + dense_output[column.name] = dense_output[column.name].replace( + to_replace=column.missing_value, method='ffill' + ) if have_sids: adjustments_from_deltas = adjustments_from_deltas_with_sids