BUG: Fixes after-hours behavior on session_date

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(

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)

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,

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)

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,