ENH: Add spot_price to BcolzDailyBarReader.

Add new method to BcolzDailyBarReader, `spot_price` which returns the
unadjusted price for the specified day and sid.
This commit is contained in:
Eddie Hebert
2015-10-10 07:19:03 -04:00
parent 4a9cd76dab
commit 5338c8e611
2 changed files with 100 additions and 0 deletions
+35
View File
@@ -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')
+65
View File
@@ -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):
"""