Merge pull request #1457 from quantopian/allow-last-date-for-session-…

…get-last-traded-dt

TST: Fix get_last_traded_dt on bcolz daily reader.

tests/data/test_resample.py

@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from unittest import skip
-
 from collections import OrderedDict
 from numbers import Real
 
@@ -666,10 +664,6 @@ def test_load_raw_arrays(self):
     def test_last_availabe_dt(self):
         self.assertEqual(self.reader.last_available_dt, self.END_DATE)
 
-    @skip("This test revealed a bug in BcolzDailyBarReader.get_last_traded_dt."
-          " When requesting data on the last session of an asset, the date is "
-          "overriden by the previous day. When that errant handling is this "
-          "test should be enabled.")
     def test_get_last_traded_dt(self):
         asset = self.asset_finder.retrieve_asset(1)
         self.assertEqual(self.reader.get_last_traded_dt(asset,

tests/data/test_us_equity_pricing.py

@@ -30,7 +30,8 @@
 
 from zipline.data.us_equity_pricing import (
     BcolzDailyBarReader,
-    NoDataOnDate,
+    NoDataBeforeDate,
+    NoDataAfterDate,
 )
 from zipline.pipeline.loaders.synthetic import (
     OHLCV,
@@ -316,11 +317,11 @@ def test_unadjusted_get_value_no_data(self):
         table = self.bcolz_daily_bar_ctable
         reader = BcolzDailyBarReader(table)
         # before
-        with self.assertRaises(NoDataOnDate):
+        with self.assertRaises(NoDataBeforeDate):
             reader.get_value(2, Timestamp('2015-06-08', tz='UTC'), 'close')
 
         # after
-        with self.assertRaises(NoDataOnDate):
+        with self.assertRaises(NoDataAfterDate):
             reader.get_value(4, Timestamp('2015-06-16', tz='UTC'), 'close')
 
     def test_unadjusted_get_value_empty_value(self):

tests/test_bar_data.py

@@ -15,6 +15,8 @@
 from datetime import timedelta
 from nose_parameterized import parameterized
 import numpy as np
+from numpy import nan
+from numpy.testing import assert_almost_equal
 import pandas as pd
 
 from zipline._protocol import handle_non_market_minutes
@@ -600,7 +602,7 @@ class TestDailyBarData(WithBarDataChecks,
                        ZiplineTestCase):
     START_DATE = pd.Timestamp('2016-01-05', tz='UTC')
     END_DATE = ASSET_FINDER_EQUITY_END_DATE = pd.Timestamp(
-        '2016-01-08',
+        '2016-01-11',
         tz='UTC',
     )
 
@@ -613,6 +615,12 @@ class TestDailyBarData(WithBarDataChecks,
     DIVIDEND_ASSET_SID = 7
     ILLIQUID_DIVIDEND_ASSET_SID = 8
 
+    @classmethod
+    def make_equity_info(cls):
+        frame = super(TestDailyBarData, cls).make_equity_info()
+        frame.loc[[1, 2], 'end_date'] = pd.Timestamp('2016-01-08', tz='UTC')
+        return frame
+
     @classmethod
     def make_splits_data(cls):
         return pd.DataFrame.from_records([
@@ -688,10 +696,11 @@ def make_adjustment_writer_equity_daily_bar_reader(cls):
     @classmethod
     def make_equity_daily_bar_data(cls):
         for sid in cls.sids:
+            asset = cls.asset_finder.retrieve_asset(sid)
             yield sid, create_daily_df_for_asset(
                 cls.trading_calendar,
-                cls.equity_daily_bar_days[0],
-                cls.equity_daily_bar_days[-1],
+                asset.start_date,
+                asset.end_date,
                 interval=2 - sid % 2
             )
 
@@ -829,25 +838,31 @@ def test_last_active_day(self):
         self.check_internal_consistency(bar_data)
 
         for asset in self.ASSETS:
-            self.assertTrue(bar_data.can_trade(asset))
+            if asset in (1, 2):
+                self.assertFalse(bar_data.can_trade(asset))
+            else:
+                self.assertTrue(bar_data.can_trade(asset))
             self.assertFalse(bar_data.is_stale(asset))
 
-            self.assertEqual(6, bar_data.current(asset, "open"))
-            self.assertEqual(7, bar_data.current(asset, "high"))
-            self.assertEqual(4, bar_data.current(asset, "low"))
-            self.assertEqual(5, bar_data.current(asset, "close"))
-            self.assertEqual(500, bar_data.current(asset, "volume"))
-            self.assertEqual(5, bar_data.current(asset, "price"))
+            if asset in (1, 2):
+                assert_almost_equal(nan, bar_data.current(asset, "open"))
+                assert_almost_equal(nan, bar_data.current(asset, "high"))
+                assert_almost_equal(nan, bar_data.current(asset, "low"))
+                assert_almost_equal(nan, bar_data.current(asset, "close"))
+                assert_almost_equal(0, bar_data.current(asset, "volume"))
+                assert_almost_equal(nan, bar_data.current(asset, "price"))
+            else:
+                self.assertEqual(6, bar_data.current(asset, "open"))
+                self.assertEqual(7, bar_data.current(asset, "high"))
+                self.assertEqual(4, bar_data.current(asset, "low"))
+                self.assertEqual(5, bar_data.current(asset, "close"))
+                self.assertEqual(500, bar_data.current(asset, "volume"))
+                self.assertEqual(5, bar_data.current(asset, "price"))
 
     def test_after_assets_dead(self):
-        # both assets end on self.day[-1], so let's try the next day
-        minute = self.get_last_minute_of_session(
-            self.trading_calendar.next_session_label(
-                self.equity_daily_bar_days[-1]
-            )
-        )
+        session = self.END_DATE
 
-        bar_data = BarData(self.data_portal, lambda: minute, "daily")
+        bar_data = BarData(self.data_portal, lambda: session, "daily")
         self.check_internal_consistency(bar_data)
 
         for asset in self.ASSETS:
@@ -861,11 +876,9 @@ def test_after_assets_dead(self):
 
             last_traded_dt = bar_data.current(asset, "last_traded")
 
-            if asset == self.ASSET1:
-                self.assertEqual(self.equity_daily_bar_days[-2],
+            if asset in (self.ASSET1, self.ASSET2):
+                self.assertEqual(self.equity_daily_bar_days[3],
                                  last_traded_dt)
-            else:
-                self.assertEqual(self.equity_daily_bar_days[1], last_traded_dt)
 
     @parameterized.expand([
         ("split", 2, 3, 3, 1.5),

zipline/data/us_equity_pricing.py

@@ -106,6 +106,14 @@ class NoDataOnDate(Exception):
     pass
 
 
+class NoDataBeforeDate(Exception):
+    pass
+
+
+class NoDataAfterDate(Exception):
+    pass
+
+
 def check_uint32_safe(value, colname):
     if value >= UINT32_MAX:
         raise ValueError(
@@ -631,19 +639,20 @@ def _spot_col(self, colname):
     def get_last_traded_dt(self, asset, day):
         volumes = self._spot_col('volume')
 
-        if day >= asset.end_date:
-            # go back to one day before the asset ended
-            search_day = self.sessions[
-                self.sessions.searchsorted(asset.end_date) - 1
-            ]
-        else:
-            search_day = day
+        search_day = day
 
         while True:
             try:
                 ix = self.sid_day_index(asset, search_day)
             except NoDataOnDate:
                 return None
+            except NoDataBeforeDate:
+                return None
+            except NoDataAfterDate:
+                prev_day_ix = self.sessions.get_loc(search_day) - 1
+                if prev_day_ix > -1:
+                    search_day = self.sessions[prev_day_ix]
+                continue
             if volumes[ix] != 0:
                 return search_day
             prev_day_ix = self.sessions.get_loc(search_day) - 1
@@ -675,12 +684,12 @@ def sid_day_index(self, sid, day):
                 day, self.sessions))
         offset = day_loc - self._calendar_offsets[sid]
         if offset < 0:
-            raise NoDataOnDate(
+            raise NoDataBeforeDate(
                 "No data on or before day={0} for sid={1}".format(
                     day, sid))
         ix = self._first_rows[sid] + offset
         if ix > self._last_rows[sid]:
-            raise NoDataOnDate(
+            raise NoDataAfterDate(
                 "No data on or after day={0} for sid={1}".format(
                     day, sid))
         return ix