Merge pull request #1153 from quantopian/filter-nulls-in-expected-cols

Filter nulls in expected cols
This commit is contained in:
Maya Tydykov
2016-04-25 16:32:45 -04:00
10 changed files with 159 additions and 118 deletions
+2 -1
View File
@@ -42,7 +42,8 @@ from zipline.testing.fixtures import (
)
date_intervals = [
[[None, '2014-01-04'], ['2014-01-05', '2014-01-09'], ['2014-01-10', None]]
[['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-09'],
['2014-01-10', '2014-01-31']]
]
buyback_authorizations_cases = [
+23 -23
View File
@@ -98,41 +98,41 @@ dividends_cases = [
prev_date_intervals = [
[
[None, '2014-01-14'], ['2014-01-15', '2014-01-19'],
['2014-01-20', None]
['2014-01-01', '2014-01-14'], ['2014-01-15', '2014-01-19'],
['2014-01-20', '2014-01-31']
],
[
[None, '2014-01-14'], ['2014-01-15', '2014-01-19'],
['2014-01-20', None]
['2014-01-01', '2014-01-14'], ['2014-01-15', '2014-01-19'],
['2014-01-20', '2014-01-31']
],
[
[None, '2014-01-09'], ['2014-01-10', '2014-01-19'],
['2014-01-20', None]
['2014-01-01', '2014-01-09'], ['2014-01-10', '2014-01-19'],
['2014-01-20', '2014-01-31']
],
[
[None, '2014-01-09'], ['2014-01-10', '2014-01-14'],
['2014-01-15', None]
['2014-01-01', '2014-01-09'], ['2014-01-10', '2014-01-14'],
['2014-01-15', '2014-01-31']
]
]
next_date_intervals = [
[
[None, '2014-01-04'], ['2014-01-05', '2014-01-15'],
['2014-01-16', '2014-01-20'], ['2014-01-21', None]
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-15'],
['2014-01-16', '2014-01-20'], ['2014-01-21', '2014-01-31']
],
[
[None, '2014-01-04'], ['2014-01-05', '2014-01-09'],
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-09'],
['2014-01-10', '2014-01-15'], ['2014-01-16', '2014-01-20'],
['2014-01-21', None]
['2014-01-21', '2014-01-31']
],
[
[None, '2014-01-04'], ['2014-01-05', '2014-01-10'],
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-10'],
['2014-01-11', '2014-01-14'], ['2014-01-15', '2014-01-20'],
['2014-01-21', None]
['2014-01-21', '2014-01-31']
],
[
[None, '2014-01-04'], ['2014-01-05', '2014-01-10'],
['2014-01-11', '2014-01-15'], ['2014-01-16', None]
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-10'],
['2014-01-11', '2014-01-15'], ['2014-01-16', '2014-01-31']
]
]
@@ -184,19 +184,19 @@ class DividendsByAnnouncementDateTestCase(WithPipelineEventDataLoader,
def setup(self, dates):
date_intervals = [
[
[None, '2014-01-04'], ['2014-01-05', '2014-01-09'],
['2014-01-10', None]
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-09'],
['2014-01-10', '2014-01-31']
],
[
[None, '2014-01-04'], ['2014-01-05', '2014-01-09'],
['2014-01-10', None]
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-09'],
['2014-01-10', '2014-01-31']
],
[
[None, '2014-01-04'], ['2014-01-05', '2014-01-14'],
['2014-01-15', None]
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-14'],
['2014-01-15', '2014-01-31']
],
[
[None, '2014-01-04'], ['2014-01-05', None]
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-31']
]
]
announcement_dates = [['NaT', '2014-01-04', '2014-01-09'],
+60 -9
View File
@@ -6,6 +6,8 @@ from unittest import TestCase
import blaze as bz
from nose_parameterized import parameterized
import numpy as np
from numpy.testing import assert_array_equal
import pandas as pd
from pandas.util.testing import assert_series_equal
@@ -28,18 +30,22 @@ from zipline.pipeline.loaders.events import (
from zipline.utils.memoize import lazyval
from zipline.utils.numpy_utils import datetime64ns_dtype
OTHER_FIELD = "other_field"
ABSTRACT_CONCRETE_LOADER_ERROR = 'abstract methods concrete_loader'
ABSTRACT_EXPECTED_COLS_ERROR = 'abstract methods expected_cols'
ABSTRACT_EXPECTED_COLS_ERROR = 'abstract methods event_date_col, expected_cols'
class EventDataSet(DataSet):
previous_announcement = Column(datetime64ns_dtype)
next_announcement = Column(datetime64ns_dtype)
class EventDataSetLoader(EventsLoader):
expected_cols = frozenset([ANNOUNCEMENT_FIELD_NAME])
event_date_col = ANNOUNCEMENT_FIELD_NAME
def __init__(self,
all_dates,
events_by_sid,
@@ -56,21 +62,21 @@ class EventDataSetLoader(EventsLoader):
def previous_announcement_loader(self):
return self._previous_event_date_loader(
self.dataset.previous_announcement,
ANNOUNCEMENT_FIELD_NAME,
)
@lazyval
def next_announcement_loader(self):
return self._previous_event_date_loader(
self.dataset.previous_announcement,
ANNOUNCEMENT_FIELD_NAME,
return self._next_event_date_loader(
self.dataset.next_announcement,
)
# Test case just for catching an error when multiple columns are in the wrong
# data format, so no loader defined.
class EventDataSetLoaderMultipleExpectedCols(EventsLoader):
expected_cols = frozenset([ANNOUNCEMENT_FIELD_NAME, "other_field"])
class EventDataSetLoaderMultipleExpectedColsNoColumnLoaders(EventsLoader):
expected_cols = frozenset([ANNOUNCEMENT_FIELD_NAME, OTHER_FIELD])
event_date_col = ANNOUNCEMENT_FIELD_NAME
class EventDataSetLoaderNoExpectedCols(EventsLoader):
@@ -92,6 +98,34 @@ dtx = pd.date_range('2014-01-01', '2014-01-10')
class EventLoaderTestCase(TestCase):
def test_null_in_event_date_col(self):
# Tests that if there is a null date in the event date column, it is
# filtered out and does not break on loading the adjusted array.
dates_with_null = pd.Series(dtx)
dates_with_null[2] = pd.NaT
events_by_sid = {0: pd.DataFrame({ANNOUNCEMENT_FIELD_NAME:
dates_with_null,
TS_FIELD_NAME: dtx})}
loader = EventDataSetLoader(
dtx,
events_by_sid,
)
prev_result = loader.load_adjusted_array({
EventDataSet.previous_announcement
}, dtx, [0], [True])[EventDataSet.previous_announcement].data[:, 0]
next_result = loader.load_adjusted_array({
EventDataSet.next_announcement
}, dtx, [0], [True])[EventDataSet.next_announcement].data[:, 0]
expected_prev = dates_with_null[:]
expected_prev[2] = dtx[1]
assert_array_equal(prev_result, expected_prev)
expected_next = dates_with_null[:]
expected_next[2] = np.datetime64('NaT')
assert_array_equal(next_result, expected_next)
def assert_loader_error(self, events_by_sid, error, msg,
infer_timestamps, loader):
with self.assertRaisesRegexp(error, re.escape(msg)):
@@ -198,7 +232,7 @@ class EventLoaderTestCase(TestCase):
[dtx, dtx],
True,
WRONG_MANY_COL_DATA_FORMAT_ERROR.format(sid=0),
EventDataSetLoaderMultipleExpectedCols
EventDataSetLoaderMultipleExpectedColsNoColumnLoaders
],
[
[dtx],
@@ -212,7 +246,7 @@ class EventLoaderTestCase(TestCase):
[dtx, dtx],
False,
WRONG_MANY_COL_DATA_FORMAT_ERROR.format(sid=0),
EventDataSetLoaderMultipleExpectedCols
EventDataSetLoaderMultipleExpectedColsNoColumnLoaders
]
]
)
@@ -246,3 +280,20 @@ class BlazeEventLoaderTestCase(TestCase):
SID_FIELD_NAME: 0})
)
)
class BlazeEventDataSetLoader(BlazeEventsLoader):
concrete_loader = EventDataSetLoader
_expected_fields = frozenset({ANNOUNCEMENT_FIELD_NAME,
TS_FIELD_NAME,
SID_FIELD_NAME})
def __init__(self,
expr,
dataset=EventDataSet,
**kwargs):
super(
BlazeEventDataSetLoader, self
).__init__(expr,
dataset=dataset,
**kwargs)