TST: add test for changing event dates and adjustments

BUG: get column names from column dict

BUG: fix name map

tests/pipeline/test_quarters_estimates.py

@@ -3,6 +3,7 @@
 from nose.tools import assert_true
 from nose_parameterized import parameterized
 import numpy as np
+from numpy.testing import assert_array_equal
 import pandas as pd
 from toolz import merge
 
@@ -657,6 +658,119 @@ def fill_expected_out(cls, expected):
         return expected
 
 
+class WithVaryingNumEstimates(WithEstimates):
+    """
+    ZiplineTestCase mixin providing fixtures and a test to ensure that we
+    have the correct overwrites when the event date changes. We want to make
+    sure that if we have a quarter with an event date that gets pushed back,
+    we don't start overwriting for the next quarter early. Likewise,
+    if we have a quarter with an event date that gets pushed forward, we want
+    to make sure that we start applying adjustments at the appropriate, earlier
+    date, rather than the later date.
+
+    Methods
+    -------
+    assert_compute()
+        Defines how to determine that results computed for the `SomeFactor`
+        factor are correct.
+
+    Tests
+    -----
+    test_windows_with_varying_num_estimates()
+        Tests that we create the correct overwrites from 2015-01-13 to
+        2015-01-14 regardless of how event dates were updated for each
+        quarter for each sid.
+    """
+
+    @classmethod
+    def make_events(cls):
+        return pd.DataFrame({
+            SID_FIELD_NAME: [0] * 3 + [1] * 3,
+            TS_FIELD_NAME: [pd.Timestamp('2015-01-09'),
+                            pd.Timestamp('2015-01-12'),
+                            pd.Timestamp('2015-01-13')] * 2,
+            EVENT_DATE_FIELD_NAME: [pd.Timestamp('2015-01-12'),
+                                    pd.Timestamp('2015-01-13'),
+                                    pd.Timestamp('2015-01-20'),
+                                    pd.Timestamp('2015-01-13'),
+                                    pd.Timestamp('2015-01-12'),
+                                    pd.Timestamp('2015-01-20')],
+            'estimate': [11., 12., 21.] * 2,
+            FISCAL_QUARTER_FIELD_NAME: [1, 1, 2] * 2,
+            FISCAL_YEAR_FIELD_NAME: [2015] * 6
+        })
+
+    @classmethod
+    def assert_compute(cls, estimate, today):
+        raise NotImplementedError('assert_compute')
+
+    def test_windows_with_varying_num_estimates(self):
+        dataset = QuartersEstimates(1)
+        assert_compute = self.assert_compute
+
+        class SomeFactor(CustomFactor):
+            inputs = [dataset.estimate]
+            window_length = 3
+
+            def compute(self, today, assets, out, estimate):
+                assert_compute(estimate, today)
+
+        engine = SimplePipelineEngine(
+            lambda x: self.loader,
+            self.trading_days,
+            self.asset_finder,
+        )
+        engine.run_pipeline(
+            Pipeline({'est': SomeFactor()}),
+            start_date=pd.Timestamp('2015-01-13', tz='utc'),
+            # last event date we have
+            end_date=pd.Timestamp('2015-01-14', tz='utc'),
+        )
+
+
+class PreviousVaryingNumEstimates(
+    WithVaryingNumEstimates,
+    ZiplineTestCase
+):
+    def assert_compute(self, estimate, today):
+        if today == pd.Timestamp('2015-01-13', tz='utc'):
+            assert_array_equal(estimate[:, 0],
+                               np.array([np.NaN, np.NaN, 12]))
+            assert_array_equal(estimate[:, 1],
+                               np.array([np.NaN, 12, 12]))
+        else:
+            assert_array_equal(estimate[:, 0],
+                               np.array([np.NaN, 12, 12]))
+            assert_array_equal(estimate[:, 1],
+                               np.array([12, 12, 12]))
+
+    @classmethod
+    def make_loader(cls, events, columns):
+        return PreviousEarningsEstimatesLoader(events, columns)
+
+
+class NextVaryingNumEstimates(
+    WithVaryingNumEstimates,
+    ZiplineTestCase
+):
+
+    def assert_compute(self, estimate, today):
+        if today == pd.Timestamp('2015-01-13', tz='utc'):
+            assert_array_equal(estimate[:, 0],
+                               np.array([11, 12, 12]))
+            assert_array_equal(estimate[:, 1],
+                               np.array([np.NaN, np.NaN, 21]))
+        else:
+            assert_array_equal(estimate[:, 0],
+                               np.array([np.NaN, 21, 21]))
+            assert_array_equal(estimate[:, 1],
+                               np.array([np.NaN, 21, 21]))
+
+    @classmethod
+    def make_loader(cls, events, columns):
+        return NextEarningsEstimatesLoader(events, columns)
+
+
 class WithEstimateWindows(WithEstimates):
     """
     ZiplineTestCase mixin providing fixures and a test to test running a

zipline/pipeline/loaders/blaze/core.py

@@ -188,7 +188,7 @@
     ensure_timezone,
     optionally,
 )
-from zipline.utils.numpy_utils import bool_dtype, categorical_dtype
+from zipline.utils.numpy_utils import bool_dtype
 from zipline.utils.pool import SequentialPool
 from zipline.utils.preprocess import preprocess
 

zipline/pipeline/loaders/blaze/estimates.py

@@ -105,20 +105,22 @@ def __init__(self,
         self._checkpoints = checkpoints
 
     def load_adjusted_array(self, columns, dates, assets, mask):
-        column_names = [column.name for column in columns]
+        # Only load requested columns.
+        requested_column_names = [self._columns[column.name]
+                                  for column in columns]
         raw = load_raw_data(
             assets,
             dates,
             self._data_query_time,
             self._data_query_tz,
-            self._expr[sorted(metadata_columns.union(column_names))],
+            self._expr[sorted(metadata_columns.union(requested_column_names))],
             self._odo_kwargs,
             checkpoints=self._checkpoints,
         )
 
         return self.loader(
             raw,
-            {k: self._columns[k] for k in column_names}
+            {column.name: self._columns[column.name] for column in columns}
         ).load_adjusted_array(
             columns,
             dates,

zipline/pipeline/loaders/blaze/events.py

@@ -58,8 +58,8 @@ class BlazeEventsLoader(PipelineLoader):
     """
 
     __doc__ = __doc__.format(SID_FIELD_NAME=SID_FIELD_NAME,
-                              TS_FIELD_NAME=TS_FIELD_NAME,
-                              EVENT_DATE_FIELD_NAME=EVENT_DATE_FIELD_NAME)
+                             TS_FIELD_NAME=TS_FIELD_NAME,
+                             EVENT_DATE_FIELD_NAME=EVENT_DATE_FIELD_NAME)
 
     @preprocess(data_query_tz=optionally(ensure_timezone))
     def __init__(self,

zipline/testing/core.py

@@ -49,14 +49,8 @@
 from zipline.utils import security_list
 from zipline.utils.calendars import get_calendar
 from zipline.utils.input_validation import expect_dimensions
-<<<<<<< HEAD
 from zipline.utils.numpy_utils import as_column, isnat
 from zipline.utils.pandas_utils import timedelta_to_integral_seconds
-=======
-from zipline.utils.numpy_utils import (
-    as_column,
-)
->>>>>>> WIP
 from zipline.utils.sentinel import sentinel
 
 import numpy as np