From 36319122cc98843ab084aebaf507166d538eaff8 Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Wed, 8 Jul 2015 23:10:37 -0400 Subject: [PATCH] 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. --- tests/test_algorithm.py | 21 +- tests/test_assets.py | 60 ++-- tests/test_perf_tracking.py | 6 +- zipline/assets/assets.py | 698 +++++++++++++++++++++++++----------- zipline/finance/trading.py | 7 - 5 files changed, 531 insertions(+), 261 deletions(-) diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 52b78ea6..f3fa3177 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -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) diff --git a/tests/test_assets.py b/tests/test_assets.py index 5c2c420d..dafa7e7a 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -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) diff --git a/tests/test_perf_tracking.py b/tests/test_perf_tracking.py index 676cd14b..ce934929 100644 --- a/tests/test_perf_tracking.py +++ b/tests/test_perf_tracking.py @@ -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), diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index 3e53a198..23784343 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -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)): diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index 0ee8085b..fe822c4e 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -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')