PERF: Change asset finder to be backed by sqlite3.

Attack the startup bottleneck of creating the asset finders caches for a
large universe, which was between 1-2 seconds on development and
production machines.

Instead, allow the AssetFinder to be passed a sqlite3 file that has
already been populated and then hydrate asset objects only when an
equity is referenced for the first time.

To create aforementioned sqlite3, create an AssetFinder with an db_path
and `create_table` set to True. If `create_table` is set to False, the
prepopulated data in the sqlite file found at db_path will be used.

Default behavior is to use an in memory database.

Behavior that changes:

- Fuzzy lookup now only works on one character, that character needs to be
specified at write/metadata consumption time, since the fuzzy lookup key
is created by dropping the character from each symbol.

- Overwriting partially written metadata is no longer
  supported. i.e. some unit tests allowed for inserting just the identifier,
  and then later updating the symbol, end_date, etc.

  Instead of building an upsert behavior at this time, this patch
  changes the unit tests so that the data for each asset is only
  inserted once.

Other notes:

- populate_cache is now removed, since there is no longer a two step
  process of inserting metadata and then realizing that metadata into
  assets. _spawn_asset is rolled into insert_metadata, so that a call to
  insert_metadata both converts the metadata and makes it available in
  the data store.
This commit is contained in:
Eddie Hebert
2015-07-14 09:54:38 -04:00
committed by Eddie Hebert
parent 3847fa70e0
commit 36319122cc
5 changed files with 531 additions and 261 deletions
+15 -6
View File
@@ -23,6 +23,7 @@ from unittest import TestCase
import numpy as np
import pandas as pd
from zipline.assets import AssetFinder
from zipline.utils.test_utils import (
nullctx,
setup_logger,
@@ -1277,16 +1278,22 @@ class TestTradingControls(TestCase):
df_source, _ = factory.create_test_df_source(self.sim_params)
metadata = {0: {'start_date': '1990-01-01',
'end_date': '2020-01-01'}}
algo = SetAssetDateBoundsAlgorithm(asset_metadata=metadata,
sim_params=self.sim_params,)
asset_finder = AssetFinder()
algo = SetAssetDateBoundsAlgorithm(
asset_finder=asset_finder,
asset_metadata=metadata,
sim_params=self.sim_params,)
algo.run(df_source)
# Run the algorithm with a sid that has already ended
df_source, _ = factory.create_test_df_source(self.sim_params)
metadata = {0: {'start_date': '1989-01-01',
'end_date': '1990-01-01'}}
algo = SetAssetDateBoundsAlgorithm(asset_metadata=metadata,
sim_params=self.sim_params,)
asset_finder = AssetFinder()
algo = SetAssetDateBoundsAlgorithm(
asset_finder=asset_finder,
asset_metadata=metadata,
sim_params=self.sim_params,)
with self.assertRaises(TradingControlViolation):
algo.run(df_source)
@@ -1294,8 +1301,10 @@ class TestTradingControls(TestCase):
df_source, _ = factory.create_test_df_source(self.sim_params)
metadata = {0: {'start_date': '2020-01-01',
'end_date': '2021-01-01'}}
algo = SetAssetDateBoundsAlgorithm(asset_metadata=metadata,
sim_params=self.sim_params,)
algo = SetAssetDateBoundsAlgorithm(
asset_finder=asset_finder,
asset_metadata=metadata,
sim_params=self.sim_params,)
with self.assertRaises(TradingControlViolation):
algo.run(df_source)
+24 -36
View File
@@ -20,14 +20,12 @@ Tests for the zipline.assets package
import sys
from unittest import TestCase
from datetime import (
timedelta,
datetime
)
from datetime import datetime, timedelta
import pickle
import uuid
import warnings
import pandas as pd
from pandas.tseries.tools import normalize_date
from nose_parameterized import parameterized
@@ -289,7 +287,7 @@ class AssetFinderTestCase(TestCase):
for i in range(3)
]
)
finder = AssetFinder(frame)
finder = AssetFinder(frame, fuzzy_char='@')
asset_0, asset_1, asset_2 = (
finder.retrieve_asset(i) for i in range(3)
)
@@ -304,17 +302,15 @@ class AssetFinderTestCase(TestCase):
# Adding an unnecessary fuzzy shouldn't matter.
self.assertEqual(
asset_1,
finder.lookup_symbol('test@1', as_of, fuzzy='@')
finder.lookup_symbol('test@1', as_of, fuzzy=True)
)
# Shouldn't find this with no fuzzy_str passed.
self.assertIsNone(finder.lookup_symbol('test1', as_of))
# Shouldn't find this with an incorrect fuzzy_str.
self.assertIsNone(finder.lookup_symbol('test1', as_of, fuzzy='*'))
# Should find it with the correct fuzzy_str.
# Should find exact match.
self.assertEqual(
asset_1,
finder.lookup_symbol('test1', as_of, fuzzy='@'),
finder.lookup_symbol('test1', as_of, fuzzy=True),
)
def test_lookup_symbol_resolve_multiple(self):
@@ -434,35 +430,28 @@ class AssetFinderTestCase(TestCase):
foo_data="FOO",)
# Test proper insertion
self.assertEqual('equity', finder.metadata_cache[0]['asset_type'])
self.assertEqual('PLAY', finder.metadata_cache[0]['symbol'])
self.assertEqual('2015-01-01', finder.metadata_cache[0]['end_date'])
equity = finder.retrieve_asset(0)
self.assertIsInstance(equity, Equity)
self.assertEqual('PLAY', equity.symbol)
self.assertEqual(pd.Timestamp('2015-01-01', tz='UTC'),
equity.end_date)
# Test invalid field
self.assertFalse('foo_data' in finder.metadata_cache[0])
# Test updating fields
finder.insert_metadata(0,
asset_type='equity',
start_date='2014-01-01',
end_date='2015-02-01',
symbol="PLAY",
exchange="NYSE",)
self.assertEqual('2015-02-01', finder.metadata_cache[0]['end_date'])
self.assertEqual('NYSE', finder.metadata_cache[0]['exchange'])
# Check that old data survived
self.assertEqual('PLAY', finder.metadata_cache[0]['symbol'])
def test_consume_metadata(self):
# Test dict consumption
finder = AssetFinder({0: {'asset_type': 'equity'}})
finder = AssetFinder()
dict_to_consume = {0: {'symbol': 'PLAY'},
1: {'symbol': 'MSFT'}}
finder.consume_metadata(dict_to_consume)
self.assertEqual('equity', finder.metadata_cache[0]['asset_type'])
self.assertEqual('PLAY', finder.metadata_cache[0]['symbol'])
equity = finder.retrieve_asset(0)
self.assertIsInstance(equity, Equity)
self.assertEqual('PLAY', equity.symbol)
finder = AssetFinder()
# Test dataframe consumption
df = pd.DataFrame(columns=['asset_name', 'exchange'], index=[0, 1])
@@ -473,11 +462,8 @@ class AssetFinderTestCase(TestCase):
finder.consume_metadata(df)
self.assertEqual('NASDAQ', finder.metadata_cache[0]['exchange'])
self.assertEqual('Microsoft', finder.metadata_cache[1]['asset_name'])
# Check that old data survived
self.assertEqual('equity', finder.metadata_cache[0]['asset_type'])
def test_consume_asset_as_identifier(self):
# Build some end dates
eq_end = pd.Timestamp('2012-01-01', tz='UTC')
fut_end = pd.Timestamp('2008-01-01', tz='UTC')
@@ -489,7 +475,6 @@ class AssetFinderTestCase(TestCase):
# Consume the Assets
finder = AssetFinder()
finder.consume_identifiers([equity_asset, future_asset])
finder.populate_cache()
# Test equality with newly built Assets
self.assertEqual(equity_asset, finder.retrieve_asset(1))
@@ -503,12 +488,15 @@ class AssetFinderTestCase(TestCase):
metadata = {'PLAY': {'symbol': 'PLAY'},
'MSFT': {'symbol': 'MSFT'}}
today = normalize_date(pd.Timestamp('2015-07-09', tz='UTC'))
# Build a finder that is allowed to assign sids
finder = AssetFinder(metadata=metadata, allow_sid_assignment=True)
finder = AssetFinder(metadata=metadata,
allow_sid_assignment=True)
# Verify that Assets were built and different sids were assigned
play = finder.lookup_symbol('PLAY', datetime.now())
msft = finder.lookup_symbol('MSFT', datetime.now())
play = finder.lookup_symbol('PLAY', today)
msft = finder.lookup_symbol('MSFT', today)
self.assertEqual('PLAY', play.symbol)
self.assertIsNotNone(play.sid)
self.assertNotEqual(play.sid, msft.sid)
+5 -1
View File
@@ -34,6 +34,7 @@ import pandas as pd
import numpy as np
from six.moves import range, zip
from zipline.assets import AssetFinder
import zipline.utils.factory as factory
import zipline.finance.performance as perf
from zipline.finance.slippage import Transaction, create_transaction
@@ -2132,7 +2133,10 @@ class TestPositionTracker(unittest.TestCase):
metadata = {1: {'asset_type': 'equity'},
2: {'asset_type': 'future',
'contract_multiplier': 1000}}
env.update_asset_finder(asset_metadata=metadata)
asset_finder = AssetFinder()
env.update_asset_finder(
asset_finder=asset_finder,
asset_metadata=metadata)
pt = perf.PositionTracker()
dt = pd.Timestamp("1984/03/06 3:00PM")
pos1 = perf.Position(1, amount=np.float64(100.0),
+487 -211
View File
@@ -14,10 +14,10 @@
# limitations under the License.
from abc import ABCMeta
from itertools import chain
from numbers import Integral
import numpy as np
import operator
import sqlite3
from sqlite3 import Row
import warnings
from logbook import Logger
@@ -63,32 +63,175 @@ ASSET_FIELDS = [
]
# Expected fields for an Asset's metadata
ASSET_TABLE_FIELDS = [
'sid',
'symbol',
'asset_name',
'start_date',
'end_date',
'first_traded',
'exchange',
]
# Expected fields for an Asset's metadata
FUTURE_TABLE_FIELDS = ASSET_TABLE_FIELDS + [
'root_symbol',
'notice_date',
'expiration_date',
'contract_multiplier',
]
EQUITY_TABLE_FIELDS = ASSET_TABLE_FIELDS
# Create the query once from the fields, so that the join is not done
# repeatedly.
FUTURE_BY_SID_QUERY = 'select {0} from futures where sid=?'.format(
", ".join(FUTURE_TABLE_FIELDS))
EQUITY_BY_SID_QUERY = 'select {0} from equities where sid=?'.format(
", ".join(EQUITY_TABLE_FIELDS))
class AssetFinder(object):
def __init__(self, metadata=None, allow_sid_assignment=True):
def __init__(self,
metadata=None,
allow_sid_assignment=True,
fuzzy_char=None,
db_path=':memory:',
create_table=True):
self.cache = {}
self.sym_cache = {}
self.future_chains_cache = {}
self.fuzzy_match = {}
self.fuzzy_char = fuzzy_char
# This flag controls if the AssetFinder is allowed to generate its own
# sids. If False, metadata that does not contain a sid will raise an
# exception when building assets.
self.allow_sid_assignment = allow_sid_assignment
# The AssetFinder also holds a nested-dict of all metadata for
# reference when building Assets
self.metadata_cache = {}
if metadata is not None:
self.consume_metadata(metadata)
if allow_sid_assignment:
self.end_date_to_assign = normalize_date(
pd.Timestamp('now', tz='UTC'))
self.populate_cache()
self.conn = sqlite3.connect(db_path)
self.cursor = self.conn.cursor()
# Create table and read in metadata.
# Should we use flags like 'r', 'w', instead?
# What we need to support is:
# - A 'throwaway' mode where the metadata is read each run.
# - A 'write' mode where the data is written to the provided db_path
# - A 'read' mode where the asset finder uses a prexisting db.
if create_table:
self.create_db_tables()
# The AssetFinder also holds a nested-dict of all metadata for
# reference when building Assets
self.metadata_cache = {}
if metadata is not None:
self.consume_metadata(metadata)
# Cache for lookup of assets by sid, the objects in the asset lookp may
# be shared with the results from equity and future lookup caches.
#
# The top level cache exists to minimize lookups on the asset type
# routing.
#
# The caches are read through, i.e. accessing an asset through
# retrieve_asset, _retrieve_equity etc. will populate the cache on
# first retrieval.
self._asset_cache = {}
self._equity_cache = {}
self._future_cache = {}
self._asset_type_cache = {}
def create_db_tables(self):
c = self.conn.cursor()
c.execute("""
CREATE TABLE equities(
sid integer,
symbol text,
asset_name text,
start_date integer,
end_date integer,
first_traded integer,
exchange text,
fuzzy text
)""")
c.execute('CREATE INDEX equities_sid on equities(sid)')
c.execute('CREATE INDEX equities_symbol on equities(symbol)')
c.execute('CREATE INDEX equities_fuzzy on equities(fuzzy)')
c.execute("""
CREATE TABLE futures(
sid integer,
symbol text,
asset_name text,
start_date integer,
end_date integer,
first_traded integer,
exchange text,
root_symbol text,
notice_date integer,
expiration_date integer,
contract_multiplier real
)""")
c.execute('CREATE INDEX futures_sid on futures(sid)')
c.execute('CREATE INDEX futures_root_symbol on equities(symbol)')
c.execute("""
CREATE TABLE asset_router
(sid integer,
asset_type text)
""")
c.execute('CREATE INDEX asset_router_sid on asset_router(sid)')
self.conn.commit()
def asset_type_by_sid(self, sid):
try:
return self._asset_type_cache[sid]
except KeyError:
pass
c = self.conn.cursor()
# Python 3 compatibility required forcing to int for sid = 0.
t = (int(sid),)
query = 'select asset_type from asset_router where sid=:sid'
c.execute(query, t)
data = c.fetchone()
if data is None:
return
asset_type = data[0]
self._asset_type_cache[sid] = asset_type
return asset_type
def retrieve_asset(self, sid, default_none=False):
if isinstance(sid, Asset):
return sid
asset = self.cache.get(sid)
try:
asset = self._asset_cache[sid]
except KeyError:
asset_type = self.asset_type_by_sid(sid)
if asset_type == 'equity':
asset = self._retrieve_equity(sid)
elif asset_type == 'future':
asset = self._retrieve_futures_contract(sid)
else:
asset = None
self._asset_cache[sid] = asset
if asset is not None:
return asset
elif default_none:
@@ -96,47 +239,71 @@ class AssetFinder(object):
else:
raise SidNotFound(sid=sid)
@staticmethod
def _lookup_symbol_in_infos(infos, as_of_date):
"""
Search a list of symbols matching a given asset for the most recent
known symbol as of as_of_date.
def _retrieve_equity(self, sid):
try:
return self._equity_cache[sid]
except KeyError:
pass
Returns a pair of (Asset, bool), representing the best match we
found for as_of_date, and whether or not that match was actually
trading at as_of_date.
c = self.conn.cursor()
c.row_factory = Row
t = (int(sid),)
c.execute(EQUITY_BY_SID_QUERY, t)
data = dict(c.fetchone())
if data:
if data['start_date']:
data['start_date'] = pd.Timestamp(data['start_date'], tz='UTC')
If no entry in infos started before as_of_date, return (None, False).
"""
# Sort entries by end_date before iterating. If asset start and end
# dates were always disjoint, then we could sort by either start or
# end_date and get the same sorting.
infos = sorted(infos, key=operator.attrgetter('end_date'))
if data['end_date']:
data['end_date'] = pd.Timestamp(data['end_date'], tz='UTC')
# Find the newest asset that started before as_of_date.
candidates = [i for i in infos
if (i.start_date is None or i.start_date <= as_of_date)
and (i.end_date is None or as_of_date <= i.end_date)]
if data['first_traded']:
data['first_traded'] = pd.Timestamp(
data['first_traded'], tz='UTC')
# If one SID exists for symbol, return that symbol
if len(candidates) == 1:
return candidates[0], True
equity = Equity(**data)
else:
equity = None
# If no SID exists for symbol, return SID with the
# highest-but-not-over end_date
if len(candidates) == 0:
candidates = [i for i in infos
if i.end_date < as_of_date]
return (candidates[-1], False) if candidates else (None, False)
self._equity_cache[sid] = equity
return equity
# If multiple SIDs exist for symbol, return latest start_date with
# end_date as a tie-breaker
if len(candidates) > 1:
best_candidate = sorted(
candidates,
key=lambda x: (x.start_date, x.end_date)
)[-1]
return best_candidate, True
def _retrieve_futures_contract(self, sid):
try:
return self._future_cache[sid]
except KeyError:
pass
c = self.conn.cursor()
t = (int(sid),)
c.row_factory = Row
c.execute(FUTURE_BY_SID_QUERY, t)
data = dict(c.fetchone())
if data:
if data['start_date']:
data['start_date'] = pd.Timestamp(data['start_date'], tz='UTC')
if data['end_date']:
data['end_date'] = pd.Timestamp(data['end_date'], tz='UTC')
if data['first_traded']:
data['first_traded'] = pd.Timestamp(
data['first_traded'], tz='UTC')
if data['notice_date']:
data['notice_date'] = pd.Timestamp(
data['notice_date'], tz='UTC')
if data['expiration_date']:
data['expiration_date'] = pd.Timestamp(
data['expiration_date'], tz='UTC')
future = Future(**data)
else:
future = None
self._future_cache[sid] = future
return future
def lookup_symbol_resolve_multiple(self, symbol, as_of_date=None):
"""
@@ -149,26 +316,70 @@ class AssetFinder(object):
raises SymbolNotFound.
"""
if as_of_date is not None:
as_of_date = normalize_date(as_of_date)
as_of_date = pd.Timestamp(normalize_date(as_of_date))
c = self.conn.cursor()
if as_of_date:
# If one SID exists for symbol, return that symbol
t = (symbol, as_of_date.value, as_of_date.value)
query = ("select sid from equities "
"where symbol=? "
"and start_date<=? "
"and end_date>=?")
c.execute(query, t)
candidates = c.fetchall()
if len(candidates) == 1:
return self._retrieve_equity(candidates[0][0])
# If no SID exists for symbol, return SID with the
# highest-but-not-over end_date
if len(candidates) == 0:
t = (symbol, as_of_date.value)
query = ("select sid from equities "
"where symbol=? "
"and start_date<=? "
"order by end_date desc "
"limit 1")
c.execute(query, t)
data = c.fetchone()
if data:
return self._retrieve_equity(data[0])
# If multiple SIDs exist for symbol, return latest start_date with
# end_date as a tie-breaker
if len(candidates) > 1:
t = (symbol, as_of_date.value)
query = ("select sid from equities "
"where symbol=? " +
"and start_date<=? " +
"order by start_date desc, end_date desc " +
"limit 1")
c.execute(query, t)
data = c.fetchone()
if data:
return self._retrieve_equity(data[0])
if symbol not in self.sym_cache:
raise SymbolNotFound(symbol=symbol)
infos = self.sym_cache[symbol]
if as_of_date is None:
if len(infos) == 1:
return infos[0]
else:
t = (symbol,)
query = ("select sid from equities where symbol=?")
c.execute(query, t)
data = c.fetchall()
if len(data) == 1:
return self._retrieve_equity(data[0][0])
elif not data:
raise SymbolNotFound(symbol=symbol)
else:
raise MultipleSymbolsFound(symbol=symbol,
options=infos)
options=str(data))
# Try to find symbol matching as_of_date
asset, _ = self._lookup_symbol_in_infos(infos, as_of_date)
if asset is None:
raise SymbolNotFound(symbol=symbol)
return asset
def lookup_symbol(self, symbol, as_of_date, fuzzy=None):
def lookup_symbol(self, symbol, as_of_date, fuzzy=False):
"""
If a fuzzy string is provided, then we try various symbols based on
the provided symbol. This is to facilitate mapping from a broker's
@@ -186,38 +397,33 @@ class AssetFinder(object):
except SymbolNotFound:
return None
else:
try:
return self.fuzzy_match[(symbol, fuzzy, as_of_date)]
except KeyError:
# if symbol is CMCSA and fuzzy is '_', then
# try CMCSA, then CMCS_A, then CMC_SA, etc.
for fuzzy_symbol in chain(
(symbol,),
(symbol[:i] + fuzzy + symbol[i:]
for i in range(len(symbol) - 1, 0, -1))):
c = self.conn.cursor()
fuzzy = symbol.replace(self.fuzzy_char, '')
t = (fuzzy, as_of_date.value, as_of_date.value)
query = ("select sid from equities "
"where fuzzy=? " +
"and start_date<=? " +
"and end_date>=?")
c.execute(query, t)
candidates = c.fetchall()
infos = self.sym_cache.get(fuzzy_symbol)
if infos:
info, date_match = self._lookup_symbol_in_infos(
infos,
as_of_date,
)
# If one SID exists for symbol, return that symbol
if len(candidates) == 1:
return self._retrieve_equity(candidates[0][0])
if info is not None and date_match:
self.fuzzy_match[(symbol, fuzzy, as_of_date)] = \
info
return info
else:
self.fuzzy_match[(symbol, fuzzy, as_of_date)] = None
def _sort_future_chains(self):
""" Sort by increasing notice date the list of contracts
for each root symbol in the future cache.
"""
notice_key = operator.attrgetter('notice_date')
for root_symbol in self.future_chains_cache:
self.future_chains_cache[root_symbol].sort(key=notice_key)
# If multiple SIDs exist for symbol, return latest start_date with
# end_date as a tie-breaker
if len(candidates) > 1:
t = (symbol, as_of_date.value)
query = ("select sid from equities "
"where symbol=? " +
"and start_date<=? " +
"order by start_date desc, end_date desc" +
"limit 1")
c.execute(query, t)
data = c.fetchone()
if data:
return self._retrieve_equity(data[0])
def lookup_future_chain(self, root_symbol, as_of_date, knowledge_date):
""" Return the futures chain for a given root symbol.
@@ -247,121 +453,37 @@ class AssetFinder(object):
Raised when a future chain could not be found for the given
root symbol.
"""
try:
return [c for c in self.future_chains_cache[root_symbol]
if c.notice_date and (as_of_date < c.notice_date)
and c.start_date and (c.start_date <= knowledge_date)]
except KeyError:
raise RootSymbolNotFound(root_symbol=root_symbol)
def populate_cache(self):
"""
Populates the asset cache with all values in the assets
collection.
"""
# Wipe caches before repopulating
self.cache = {}
self.sym_cache = {}
self.future_chains_cache = {}
self.fuzzy_match = {}
for identifier, row in self.metadata_cache.items():
asset = self._spawn_asset(identifier=identifier, **row)
# Insert asset into the various caches
self.cache[asset.sid] = asset
if asset.symbol is not '':
self.sym_cache.setdefault(asset.symbol, []).append(asset)
if isinstance(asset, Future) and asset.root_symbol is not '':
self.future_chains_cache.setdefault(asset.root_symbol,
[]).append(asset)
# Pre-sort the future chains, we assume in future lookups
# that they're ordered correctly.
self._sort_future_chains()
def _spawn_asset(self, identifier, **kwargs):
# If the file_name is in the kwargs, it will be used as the symbol
try:
kwargs['symbol'] = kwargs.pop('file_name')
except KeyError:
pass
# If the identifier coming in was a string and there is no defined
# symbol yet, set the symbol to the incoming identifier
try:
kwargs['symbol']
pass
except KeyError:
if isinstance(identifier, string_types):
kwargs['symbol'] = identifier
# If the company_name is in the kwargs, it may be the asset_name
try:
company_name = kwargs.pop('company_name')
try:
kwargs['asset_name']
except KeyError:
kwargs['asset_name'] = company_name
except KeyError:
pass
# If dates are given as nanos, pop them
try:
kwargs['start_date'] = kwargs.pop('start_date_nano')
except KeyError:
pass
try:
kwargs['end_date'] = kwargs.pop('end_date_nano')
except KeyError:
pass
try:
kwargs['notice_date'] = kwargs.pop('notice_date_nano')
except KeyError:
pass
try:
kwargs['expiration_date'] = kwargs.pop('expiration_date_nano')
except KeyError:
pass
# Process dates to Timestamps
try:
kwargs['start_date'] = pd.Timestamp(kwargs['start_date'], tz='UTC')
except KeyError:
pass
try:
kwargs['end_date'] = pd.Timestamp(kwargs['end_date'], tz='UTC')
except KeyError:
pass
try:
kwargs['notice_date'] = pd.Timestamp(kwargs['notice_date'],
tz='UTC')
except KeyError:
pass
try:
kwargs['expiration_date'] = pd.Timestamp(kwargs['expiration_date'],
tz='UTC')
except KeyError:
pass
# Build an Asset of the appropriate type, default to Equity
asset_type = kwargs.pop('asset_type', 'equity')
if asset_type.lower() == 'equity':
asset = Equity(**kwargs)
elif asset_type.lower() == 'future':
asset = Future(**kwargs)
else:
raise InvalidAssetType(asset_type=asset_type)
return asset
c = self.conn.cursor()
t = {'root_symbol': root_symbol,
'as_of_date': as_of_date.value,
'knowledge_date': knowledge_date.value}
c.execute("""
select sid from futures
where root_symbol=:root_symbol
and :as_of_date < notice_date
and start_date <= :knowledge_date
order by notice_date asc
""", t)
sids = [r[0] for r in c.fetchall()]
if not sids:
# Check if root symbol exists.
c.execute("""
select count(sid) from futures where root_symbol=:root_symbol
""", t)
count = c.fetchone()[0]
if count == 0:
raise RootSymbolNotFound(root_symbol=root_symbol)
else:
# If symbol exists, return empty future chain.
return []
return [self._retrieve_futures_contract(sid) for sid in sids]
@property
def sids(self):
return self.cache.keys()
c = self.conn.cursor()
query = 'select sid from asset_router'
c.execute(query)
return [r[0] for r in c.fetchall()]
@property
def assets(self):
@@ -490,7 +612,6 @@ class AssetFinder(object):
# If symbols or Assets are provided, construction and mapping is
# necessary
self.consume_identifiers(index)
self.populate_cache()
# Look up all Assets for mapping
matches = []
@@ -506,14 +627,22 @@ class AssetFinder(object):
# Return a list of the sids of the found assets
return [asset.sid for asset in matches]
def insert_metadata(self, identifier, **kwargs):
def _insert_metadata(self, identifier, **kwargs):
"""
Inserts the given metadata kwargs to the entry for the given
identifier. Matching fields in the existing entry will be overwritten.
:param identifier: The identifier for which to insert metadata
:param kwargs: The keyed metadata to insert
"""
entry = self.metadata_cache.get(identifier, {})
if identifier in self.metadata_cache:
# Multiple pass insertion no longer supported.
# This could and probably should raise an Exception, but is
# currently just a short-circuit for compatibility with existing
# testing structure in the test_algorithm module which creates
# multiple sources which all insert redundant metadata.
return
entry = {}
for key, value in kwargs.items():
# Do not accept invalid fields
@@ -545,6 +674,140 @@ class AssetFinder(object):
else:
raise SidAssignmentError(identifier=identifier)
# If the file_name is in the kwargs, it will be used as the symbol
try:
entry['symbol'] = entry.pop('file_name')
except KeyError:
pass
# If the identifier coming in was a string and there is no defined
# symbol yet, set the symbol to the incoming identifier
try:
entry['symbol']
pass
except KeyError:
if isinstance(identifier, string_types):
entry['symbol'] = identifier
# If the company_name is in the kwargs, it may be the asset_name
try:
company_name = entry.pop('company_name')
try:
entry['asset_name']
except KeyError:
entry['asset_name'] = company_name
except KeyError:
pass
# If dates are given as nanos, pop them
try:
entry['start_date'] = entry.pop('start_date_nano')
except KeyError:
pass
try:
entry['end_date'] = entry.pop('end_date_nano')
except KeyError:
pass
try:
entry['notice_date'] = entry.pop('notice_date_nano')
except KeyError:
pass
try:
entry['expiration_date'] = entry.pop('expiration_date_nano')
except KeyError:
pass
# Process dates to Timestamps
try:
entry['start_date'] = pd.Timestamp(entry['start_date'], tz='UTC')
except KeyError:
# Set a default start_date of the EPOCH, so that all date queries
# work when a start date is not provided.
entry['start_date'] = pd.Timestamp(0, tz='UTC')
try:
# Set a default end_date of 'now', so that all date queries
# work when a end date is not provided.
entry['end_date'] = pd.Timestamp(entry['end_date'], tz='UTC')
except KeyError:
entry['end_date'] = self.end_date_to_assign
try:
entry['notice_date'] = pd.Timestamp(entry['notice_date'],
tz='UTC')
except KeyError:
pass
try:
entry['expiration_date'] = pd.Timestamp(entry['expiration_date'],
tz='UTC')
except KeyError:
pass
# Build an Asset of the appropriate type, default to Equity
asset_type = entry.pop('asset_type', 'equity')
if asset_type.lower() == 'equity':
fuzzy = entry['symbol'].replace(self.fuzzy_char, '') \
if self.fuzzy_char else None
asset = Equity(**entry)
c = self.conn.cursor()
t = (asset.sid,
asset.symbol,
asset.asset_name,
asset.start_date.value if asset.start_date else None,
asset.end_date.value if asset.end_date else None,
asset.first_traded.value if asset.first_traded else None,
asset.exchange,
fuzzy)
c.execute("""INSERT INTO equities(
sid,
symbol,
asset_name,
start_date,
end_date,
first_traded,
exchange,
fuzzy)
VALUES(?, ?, ?, ?, ?, ?, ?, ?)""", t)
t = (asset.sid,
'equity')
c.execute("""INSERT INTO asset_router(sid, asset_type)
VALUES(?, ?)""", t)
elif asset_type.lower() == 'future':
asset = Future(**entry)
c = self.conn.cursor()
t = (asset.sid,
asset.symbol,
asset.asset_name,
asset.start_date.value if asset.start_date else None,
asset.end_date.value if asset.end_date else None,
asset.first_traded.value if asset.first_traded else None,
asset.exchange,
asset.root_symbol,
asset.notice_date.value if asset.notice_date else None,
asset.expiration_date.value
if asset.expiration_date else None,
asset.contract_multiplier)
c.execute("""INSERT INTO futures(
sid,
symbol,
asset_name,
start_date,
end_date,
first_traded,
exchange,
root_symbol,
notice_date,
expiration_date,
contract_multiplier)
VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", t)
t = (asset.sid,
'future')
c.execute("""INSERT INTO asset_router(sid, asset_type)
VALUES(?, ?)""", t)
else:
raise InvalidAssetType(asset_type=asset_type)
self.metadata_cache[identifier] = entry
def consume_identifiers(self, identifiers):
@@ -584,15 +847,27 @@ class AssetFinder(object):
raise ConsumeAssetMetaDataError(obj=metadata)
def clear_metadata(self):
"""
Used for testing.
"""
self.metadata_cache = {}
self.conn = sqlite3.connect(':memory:')
self.create_db_tables()
def insert_metadata(self, identifier, **kwargs):
self._insert_metadata(identifier, **kwargs)
self.conn.commit()
def _insert_metadata_dataframe(self, dataframe):
for identifier, row in dataframe.iterrows():
self.insert_metadata(identifier, **row)
self._insert_metadata(identifier, **row)
self.conn.commit()
def _insert_metadata_dict(self, dict):
for identifier, entry in dict.items():
self.insert_metadata(identifier, **entry)
self._insert_metadata(identifier, **entry)
self.conn.commit()
def _insert_metadata_readable(self, readable):
for row in readable.read():
@@ -615,7 +890,8 @@ class AssetFinder(object):
identifier = metadata_dict['symbol']
else:
raise ConsumeAssetMetaDataError(obj=row)
self.insert_metadata(identifier, **metadata_dict)
self._insert_metadata(identifier, **metadata_dict)
self.conn.commit()
class AssetConvertible(with_metaclass(ABCMeta)):
-7
View File
@@ -177,10 +177,8 @@ class TradingEnvironment(object):
:param identifiers: A list of identifiers to be inserted
:return:
"""
populate = False
if clear_metadata:
self.asset_finder.clear_metadata()
populate = True
if asset_finder is not None:
if not isinstance(asset_finder, AssetFinder):
@@ -190,14 +188,9 @@ class TradingEnvironment(object):
if asset_metadata is not None:
self.asset_finder.clear_metadata()
self.asset_finder.consume_metadata(asset_metadata)
populate = True
if identifiers is not None:
self.asset_finder.consume_identifiers(identifiers)
populate = True
if populate:
self.asset_finder.populate_cache()
def normalize_date(self, test_date):
test_date = pd.Timestamp(test_date, tz='UTC')