diff --git a/tests/test_data_portal.py b/tests/test_data_portal.py index 485c0b62..71b50bd9 100644 --- a/tests/test_data_portal.py +++ b/tests/test_data_portal.py @@ -42,7 +42,7 @@ class TestDataPortal(WithTradingEnvironment, ZiplineTestCase): # all the minutes of 7/6, 7/7, 7/8, and 31 minutes of 7/9 july_9_dt = self.trading_schedule.start_and_end( - pd.Timestamp("2015-07-09") + pd.Timestamp("2015-07-09", tz='UTC') )[0] + Timedelta("30 minutes") self.assertEqual( @@ -65,7 +65,7 @@ class TestDataPortal(WithTradingEnvironment, ZiplineTestCase): # all the minutes of 11/24, 11/25, 11/27 (half day!), and 31 minutes # of 11/30 nov_30_dt = self.trading_schedule.start_and_end( - pd.Timestamp("2015-11-30") + pd.Timestamp("2015-11-30", tz='UTC') )[0] + Timedelta("30 minutes") self.assertEqual( diff --git a/zipline/utils/calendars/calendar_helpers.py b/zipline/utils/calendars/calendar_helpers.py index e235af77..2e2f7162 100644 --- a/zipline/utils/calendars/calendar_helpers.py +++ b/zipline/utils/calendars/calendar_helpers.py @@ -49,12 +49,24 @@ def _get_index(dt, all_trading_days): return all_trading_days.searchsorted(ndt) - 1 # The following methods are intended to be inserted in both the -# ExchangeCalendar and TradingSchedule classes with partial hooks to those -# class' methods. These methods live in the helpers module to avoid code -# duplication. +# ExchangeCalendar and TradingSchedule classes. +# These methods live in the helpers module to avoid code duplication. def next_scheduled_day(date, last_trading_day, is_scheduled_day_hook): + """ + Returns the next session date in the calendar after the provided date. + + Parameters + ---------- + date : Timestamp + The date whose following date is needed. + + Returns + ------- + Timestamp + The next scheduled date after the provided date. + """ dt = normalize_date(date) delta = pd.Timedelta(days=1) @@ -66,6 +78,19 @@ def next_scheduled_day(date, last_trading_day, is_scheduled_day_hook): def previous_scheduled_day(date, first_trading_day, is_scheduled_day_hook): + """ + Returns the previous session date in the calendar before the provided date. + + Parameters + ---------- + date : Timestamp + The date whose previous date is needed. + + Returns + ------- + Timestamp + The previous scheduled date before the provided date. + """ dt = normalize_date(date) delta = pd.Timedelta(days=-1) @@ -141,16 +166,18 @@ def add_scheduled_days(n, date, next_scheduled_day_hook, Adds n trading days to date. If this would fall outside of the trading calendar, a NoFurtherDataError is raised. - :Arguments: - n : int - The number of days to add to date, this can be positive or - negative. - date : datetime - The date to add to. + Parameters + ---------- + n : int + The number of days to add to date, this can be positive or + negative. + date : datetime + The date to add to. - :Returns: - new_date : datetime - n trading days added to date. + Returns + ------- + datetime + n trading days added to date. """ if n == 1: return next_scheduled_day_hook(date) diff --git a/zipline/utils/calendars/exchange_calendar.py b/zipline/utils/calendars/exchange_calendar.py index e851ae78..5d832234 100644 --- a/zipline/utils/calendars/exchange_calendar.py +++ b/zipline/utils/calendars/exchange_calendar.py @@ -240,6 +240,23 @@ class ExchangeCalendar(with_metaclass(ABCMeta)): ) def add_trading_days(self, n, date): + """ + Adds n trading days to date. If this would fall outside of the + ExchangeCalendar, a NoFurtherDataError is raised. + + Parameters + ---------- + n : int + The number of days to add to date, this can be positive or + negative. + date : datetime + The date to add to. + + Returns + ------- + datetime + n trading days added to date. + """ return add_scheduled_days( n, date, next_scheduled_day_hook=self.next_trading_day, diff --git a/zipline/utils/calendars/nyse_exchange_calendar.py b/zipline/utils/calendars/nyse_exchange_calendar.py index 4032333a..34fd2633 100644 --- a/zipline/utils/calendars/nyse_exchange_calendar.py +++ b/zipline/utils/calendars/nyse_exchange_calendar.py @@ -437,6 +437,14 @@ class NYSEExchangeCalendar(ExchangeCalendar): Timestamp The date of the exchange session in which dt belongs. """ + # Check if the dt is after the market close + # If so, advance to the next day + if self.is_open_on_day(dt): + _, close = self._get_open_and_close(normalize_date(dt)) + if dt > close: + dt += Timedelta(days=1) + while not self.is_open_on_day(dt): dt += Timedelta(days=1) + return normalize_date(dt) diff --git a/zipline/utils/calendars/trading_schedule.py b/zipline/utils/calendars/trading_schedule.py index 0bd86aa3..3cf02cb6 100644 --- a/zipline/utils/calendars/trading_schedule.py +++ b/zipline/utils/calendars/trading_schedule.py @@ -97,6 +97,23 @@ class TradingSchedule(with_metaclass(ABCMeta)): ) def add_execution_days(self, n, date): + """ + Adds n execution days to date. If this would fall outside of the + TradingSchedule, a NoFurtherDataError is raised. + + Parameters + ---------- + n : int + The number of days to add to date, this can be positive or + negative. + date : datetime + The date to add to. + + Returns + ------- + datetime + n trading days added to date. + """ return add_scheduled_days( n, date, next_scheduled_day_hook=self.next_execution_day,