ENH: Add run_chunked_pipeline method to PipelineEngine

2026-08-11 11:16:15 +08:00 · 2017-06-02 16:48:09 -04:00
parent 69b632db37
commit 2d56d253fa
8 changed files with 191 additions and 125 deletions
@@ -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."
+        )