ENH: Add run_chunked_pipeline method to PipelineEngine

tests/pipeline/test_engine.py

@@ -51,6 +51,7 @@ from zipline.pipeline.factors import (
     ExponentialWeightedMovingAverage,
     ExponentialWeightedMovingStdDev,
     MaxDrawdown,
+    Returns,
     SimpleMovingAverage,
 )
 from zipline.pipeline.loaders.equity_pricing_loader import (
@@ -77,6 +78,7 @@ from zipline.testing import (
 )
 from zipline.testing.fixtures import (
     WithAdjustmentReader,
+    WithEquityPricingPipelineEngine,
     WithSeededRandomPipelineEngine,
     WithTradingEnvironment,
     ZiplineTestCase,
@@ -1497,3 +1499,35 @@ class PopulateInitialWorkspaceTestCase(WithConstantInputs, ZiplineTestCase):
                 precomputed_term_value,
             ),
         )
+
+
+class ChunkedPipelineTestCase(WithEquityPricingPipelineEngine,
+                              ZiplineTestCase):
+
+    PIPELINE_START_DATE = Timestamp('2006-01-05', tz='UTC')
+    END_DATE = Timestamp('2006-12-29', tz='UTC')
+
+    def test_run_chunked_pipeline(self):
+        """
+        Test that running a pipeline in chunks produces the same result as if
+        it were run all at once
+        """
+        pipe = Pipeline(
+            columns={
+                'close': USEquityPricing.close.latest,
+                'returns': Returns(window_length=2),
+                'categorical': USEquityPricing.close.latest.quantiles(5)
+            },
+        )
+        pipeline_result = self.pipeline_engine.run_pipeline(
+            pipe,
+            start_date=self.PIPELINE_START_DATE,
+            end_date=self.END_DATE,
+        )
+        chunked_result = self.pipeline_engine.run_chunked_pipeline(
+            pipeline=pipe,
+            start_date=self.PIPELINE_START_DATE,
+            end_date=self.END_DATE,
+            chunksize=22
+        )
+        self.assertTrue(chunked_result.equals(pipeline_result))

tests/pipeline/test_run_chunked_pipeline.py

@@ -1,37 +0,0 @@
-from zipline.pipeline import Pipeline, run_chunked_pipeline
-from zipline.pipeline.data import USEquityPricing
-from zipline.pipeline.factors import Returns
-from zipline.testing import ZiplineTestCase
-from zipline.testing.fixtures import WithEquityPricingPipelineEngine
-
-
-class ChunkedPipelineTestCase(WithEquityPricingPipelineEngine,
-                              ZiplineTestCase):
-
-    def test_run_chunked_pipeline(self):
-        """
-        Test that running a pipeline in chunks produces the same result as if
-        it were run all at once
-        """
-        pipe = Pipeline(
-            columns={
-                'close': USEquityPricing.close.latest,
-                'returns': Returns(window_length=2),
-            },
-        )
-        sessions = self.nyse_calendar.all_sessions
-        start_date = sessions[sessions.get_loc(self.START_DATE) + 2]
-
-        pipeline_result = self.pipeline_engine.run_pipeline(
-            pipe,
-            start_date=start_date,
-            end_date=self.END_DATE,
-        )
-        chunked_result = run_chunked_pipeline(
-            engine=self.pipeline_engine,
-            pipeline=pipe,
-            start_date=start_date,
-            end_date=self.END_DATE,
-            chunksize=22
-        )
-        self.assertTrue(chunked_result.equals(pipeline_result))

tests/utils/test_date_utils.py

@@ -3,26 +3,84 @@ from nose_parameterized import parameterized
 
 from zipline.testing import ZiplineTestCase
 from zipline.utils.calendars import get_calendar
-from zipline.utils.date_utils import roll_dates_to_previous_session
+from zipline.utils.date_utils import compute_date_range_chunks
 
 
-class TestRollDatesToPreviousSession(ZiplineTestCase):
+def T(s):
+    """
+    Helpful function to improve readibility.
+    """
+    return Timestamp(s, tz='UTC')
+
+
+class TestDateUtils(ZiplineTestCase):
+
+    @classmethod
+    def init_class_fixtures(cls):
+        super(TestDateUtils, cls).init_class_fixtures()
+        cls.calendar = get_calendar('NYSE')
 
     @parameterized.expand([
-        (
-            Timestamp('05-19-2017', tz='UTC'),  # actual trading date
-            Timestamp('05-19-2017', tz='UTC'),
-        ),
-        (
-            Timestamp('07-04-2015', tz='UTC'),  # weekend nyse holiday
-            Timestamp('07-02-2015', tz='UTC'),
-        ),
-        (
-            Timestamp('01-16-2017', tz='UTC'),  # weeknight nyse holiday
-            Timestamp('01-13-2017', tz='UTC'),
-        ),
+        (None, [(T('2017-01-03'), T('2017-01-31'))]),
+        (10, [
+            (T('2017-01-03'), T('2017-01-17')),
+            (T('2017-01-18'), T('2017-01-31'))
+        ]),
+        (15, [
+            (T('2017-01-03'), T('2017-01-24')),
+            (T('2017-01-25'), T('2017-01-31'))
+        ]),
     ])
-    def test_roll_dates_to_previous_session(self, date, expected_rolled_date):
-        calendar = get_calendar('NYSE')
-        result = roll_dates_to_previous_session(calendar, date)
-        self.assertEqual(result[0], expected_rolled_date)
+    def test_compute_date_range_chunks(self, chunksize, expected):
+        # This date range results in 20 business days
+        start_date = T('2017-01-03')
+        end_date = T('2017-01-31')
+
+        date_ranges = compute_date_range_chunks(
+            self.calendar.all_sessions,
+            start_date,
+            end_date,
+            chunksize
+        )
+
+        self.assertListEqual(list(date_ranges), expected)
+
+    def test_compute_date_range_chunks_invalid_input(self):
+        # Start date not found in calendar
+        with self.assertRaises(KeyError) as cm:
+            compute_date_range_chunks(
+                self.calendar.all_sessions,
+                T('2017-05-07'),  # Sunday
+                T('2017-06-01'),
+                None
+            )
+        self.assertEqual(
+            str(cm.exception),
+            "'Start date 2017-05-07 is not found in calendar.'"
+        )
+
+        # End date not found in calendar
+        with self.assertRaises(KeyError) as cm:
+            compute_date_range_chunks(
+                self.calendar.all_sessions,
+                T('2017-05-01'),
+                T('2017-05-27'),  # Saturday
+                None
+            )
+        self.assertEqual(
+            str(cm.exception),
+            "'End date 2017-05-27 is not found in calendar.'"
+        )
+
+        # End date before start date
+        with self.assertRaises(ValueError) as cm:
+            compute_date_range_chunks(
+                self.calendar.all_sessions,
+                T('2017-06-01'),
+                T('2017-05-01'),
+                None
+            )
+        self.assertEqual(
+            str(cm.exception),
+            "End date 2017-05-01 cannot precede start date 2017-06-01."
+        )

tests/utils/test_pandas_utils.py

@@ -173,8 +173,16 @@ class TestCatDFConcat(ZiplineTestCase):
             ),
         ]
 
-        with self.assertRaises(ValueError):
+        with self.assertRaises(ValueError) as cm:
             categorical_df_concat(mismatched_dtypes)
+        self.assertEqual(
+            str(cm.exception),
+            "Input DataFrames must have the same columns/dtypes."
+        )
 
-        with self.assertRaises(ValueError):
+        with self.assertRaises(ValueError) as cm:
             categorical_df_concat(mismatched_column_names)
+        self.assertEqual(
+            str(cm.exception),
+            "Input DataFrames must have the same columns/dtypes."
+        )