mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-29 11:18:20 +08:00
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:
@@ -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')
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user