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')