From 5338c8e61121dc4995b30520f3be9ba19c725b96 Mon Sep 17 00:00:00 2001
From: Eddie Hebert <ehebert@quantopian.com>
Date: Sat, 10 Oct 2015 00:02:23 -0400
Subject: [PATCH] ENH: Add spot_price to BcolzDailyBarReader.

Add new method to BcolzDailyBarReader, `spot_price` which returns the
unadjusted price for the specified day and sid.
---
 tests/data/test_us_equity_pricing.py | 35 +++++++++++++++
 zipline/data/us_equity_pricing.py    | 65 ++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+)

diff --git a/tests/data/test_us_equity_pricing.py b/tests/data/test_us_equity_pricing.py
index f0a418c5..a6431875 100644
--- a/tests/data/test_us_equity_pricing.py
+++ b/tests/data/test_us_equity_pricing.py
@@ -35,6 +35,7 @@ from zipline.pipeline.loaders.synthetic import (
 )
 from zipline.data.us_equity_pricing import (
     BcolzDailyBarReader,
+    NoDataOnDate
 )
 from zipline.finance.trading import TradingEnvironment
 from zipline.pipeline.data import USEquityPricing
@@ -266,3 +267,37 @@ class BcolzDailyBarTestCase(TestCase):
                 start_date=self.trading_days[0],
                 end_date=self.asset_end(asset),
             )
+
+    def test_unadjusted_spot_price(self):
+        table = self.writer.write(self.dest, self.trading_days, self.assets)
+        reader = BcolzDailyBarReader(table)
+        # At beginning
+        price = reader.spot_price(1, Timestamp('2015-06-01', tz='UTC'),
+                                  'close')
+        # Synthetic writes price for date.
+        self.assertEqual(135630.0, price)
+
+        # Middle
+        price = reader.spot_price(1, Timestamp('2015-06-02', tz='UTC'),
+                                  'close')
+        self.assertEqual(135631.0, price)
+        # End
+        price = reader.spot_price(1, Timestamp('2015-06-05', tz='UTC'),
+                                  'close')
+        self.assertEqual(135634.0, price)
+
+        # Another sid at beginning.
+        price = reader.spot_price(2, Timestamp('2015-06-22', tz='UTC'),
+                                  'close')
+        self.assertEqual(235651.0, price)
+
+    def test_unadjusted_spot_price_no_data(self):
+        table = self.writer.write(self.dest, self.trading_days, self.assets)
+        reader = BcolzDailyBarReader(table)
+        # before
+        with self.assertRaises(NoDataOnDate):
+            reader.spot_price(2, Timestamp('2015-06-08', tz='UTC'), 'close')
+
+        # after
+        with self.assertRaises(NoDataOnDate):
+            reader.spot_price(4, Timestamp('2015-06-16', tz='UTC'), 'close')
diff --git a/zipline/data/us_equity_pricing.py b/zipline/data/us_equity_pricing.py
index 302e1c75..4cac87d6 100644
--- a/zipline/data/us_equity_pricing.py
+++ b/zipline/data/us_equity_pricing.py
@@ -64,6 +64,13 @@ SQLITE_ADJUSTMENT_TABLENAMES = frozenset(['splits', 'dividends', 'mergers'])
 UINT32_MAX = iinfo(uint32).max
 
 
+class NoDataOnDate(Exception):
+    """
+    Raised when a spot price can be found for the sid and date.
+    """
+    pass
+
+
 class BcolzDailyBarWriter(with_metaclass(ABCMeta)):
     """
     Class capable of writing daily OHLCV data to disk in a format that can be
@@ -333,6 +340,11 @@ class BcolzDailyBarReader(object):
             int(id_): offset
             for id_, offset in iteritems(table.attrs['calendar_offset'])
         }
+        # Cache of fully read np.array for the carrays in the daily bar table.
+        # raw_array does not use the same cache, but it could.
+        # Need to test keeping the entire array in memory for the course of a
+        # process first.
+        self._spot_cols = {}
 
     def _compute_slices(self, start_idx, end_idx, assets):
         """
@@ -394,6 +406,59 @@ class BcolzDailyBarReader(object):
             offsets,
         )
 
+    def _spot_col(self, colname):
+        """
+        Get the colname from daily_bar_table and read all of it into memory,
+        caching the result.
+
+        Parameters
+        ----------
+        colname : string
+            A name of a OHLCV carray in the daily_bar_table
+
+        Returns
+        -------
+        array (uint32)
+            Full read array of the carray in the daily_bar_table with the
+            given colname.
+        """
+        try:
+            col = self._spot_cols[colname]
+        except KeyError:
+            col = self._spot_cols[colname] = self._table[colname][:]
+        return col
+
+    def spot_price(self, sid, day, colname):
+        """
+        Parameters
+        ----------
+        sid : int
+            The asset identifier.
+        day : datetime64
+            Midnight of the day for which data is requested.
+        colname : string
+            The price field. e.g. ('open', 'high', 'low', 'close', 'volume')
+
+        Returns
+        -------
+        float
+            The spot price for colname of the given sid on the given day.
+            Raises a NoDataOnDate exception if there is no data available
+            for the given day and sid.
+        """
+        day_loc = self._calendar.get_loc(day)
+        offset = day_loc - self._calendar_offsets[sid]
+        if offset < 0:
+            raise NoDataOnDate(
+                "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(
+                "No data on or after day={0} for sid={1}".format(
+                    day, sid))
+        return self._spot_col(colname)[ix] * 0.001
+
 
 class SQLiteAdjustmentWriter(object):
     """