diff --git a/tests/test_assets.py b/tests/test_assets.py index de2561d2..fa53794f 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -446,7 +446,8 @@ class AssetFinderTestCase(TestCase): equity.end_date) # Test invalid field - self.assertFalse('foo_data' in finder.metadata_cache[0]) + with self.assertRaises(AttributeError): + equity.foo_data def test_consume_metadata(self): @@ -469,8 +470,8 @@ class AssetFinderTestCase(TestCase): df['asset_name'][1] = "Microsoft" df['exchange'][1] = "NYSE" finder.consume_metadata(df) - self.assertEqual('NASDAQ', finder.metadata_cache[0]['exchange']) - self.assertEqual('Microsoft', finder.metadata_cache[1]['asset_name']) + self.assertEqual('NASDAQ', finder.retrieve_asset(0).exchange) + self.assertEqual('Microsoft', finder.retrieve_asset(1).asset_name) def test_consume_asset_as_identifier(self): # Build some end dates diff --git a/zipline/assets/__init__.py b/zipline/assets/__init__.py index b088688e..0363ffde 100644 --- a/zipline/assets/__init__.py +++ b/zipline/assets/__init__.py @@ -20,6 +20,14 @@ from ._assets import ( make_asset_array, CACHE_FILE_TEMPLATE ) +from .asset_writer import ( + AssetDBWriterFromDictionary, + NullAssetDBWriterLegacy, + AssetDBWriterLegacyFromList, + AssetDBWriterLegacyFromDictionary, + AssetDBWriterLegacyFromDataFrame, + AssetDBWriterLegacyFromReadable +) from .assets import ( AssetFinder, AssetConvertible @@ -29,6 +37,12 @@ __all__ = [ 'Asset', 'Equity', 'Future', + 'AssetDBWriterFromDictionary', + 'NullAssetDBWriterLegacy', + 'AssetDBWriterLegacyFromList', + 'AssetDBWriterLegacyFromDictionary', + 'AssetDBWriterLegacyFromDataFrame', + 'AssetDBWriterLegacyFromReadable', 'AssetFinder', 'AssetConvertible', 'make_asset_array', diff --git a/zipline/assets/_assets.pyx b/zipline/assets/_assets.pyx index f71f86b1..fe628f2f 100644 --- a/zipline/assets/_assets.pyx +++ b/zipline/assets/_assets.pyx @@ -232,6 +232,7 @@ cdef class Future(Asset): def __cinit__(self, int sid, # sid is required object symbol="", + object root_symbol_id = "", object root_symbol="", object asset_name="", object start_date=None, diff --git a/zipline/assets/asset_writer.py b/zipline/assets/asset_writer.py new file mode 100644 index 00000000..e529a660 --- /dev/null +++ b/zipline/assets/asset_writer.py @@ -0,0 +1,711 @@ +import pandas as pd +import numpy as np +from pandas.tseries.tools import normalize_date +from six import with_metaclass, string_types +from abc import ( + ABCMeta, + abstractmethod, +) +from zipline.errors import ( + ConsumeAssetMetaDataError, + InvalidAssetType, + SidAssignmentError, +) +from zipline.assets import ( + Asset, Equity, Future +) + +ASSET_FIELDS = [ + 'sid', + 'asset_type', + 'symbol', + 'asset_name', + 'start_date', + 'end_date', + 'first_traded', + 'exchange', + 'notice_date', + 'root_symbol', + 'expiration_date', + 'contract_multiplier', + # The following fields are for compatibility with other systems + 'file_name', # Used as symbol + 'company_name', # Used as asset_name + 'start_date_nano', # Used as start_date + 'end_date_nano' # Used as end_date +] + +# 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_id', + 'notice_date', + 'expiration_date', + 'contract_multiplier', +] + +EQUITY_TABLE_FIELDS = ASSET_TABLE_FIELDS + +EXCHANGE_TABLE_FIELDS = [ + 'exchange_id', + 'exchange', + 'timezone' +] + +ROOT_SYMBOL_TABLE_FIELDS = [ + 'root_symbol_id', + 'root_symbol', + 'sector', + 'description', + 'exchange_id' +] + + +class AssetDBWriter(with_metaclass(ABCMeta)): + """ + Class used to write arbitrary data to SQLite database. + Concrete subclasses will implement the logic for a specific + input datatypes by implementing the load_data method. + + Methods + ------- + write_all(db_conn, fuzzy_char=None, allow_sid_assignment=True, + constraints=False) + Write the data supplied at initialization to the database. + init_db(db_conn, constraints=False) + Create the SQLite tables (called by write_all). + load_data(self) + Returns data in standard format. + + """ + + def write_all(self, db_conn, fuzzy_char=None, allow_sid_assignment=True, + constraints=False): + """ Write pre-supplied data to SQLite. + + Parameters + ---------- + db_conn: sqlite3.Connection + A connection to a SQLite database. + fuzzy_char: string + A string for use in fuzzy matching. + allow_sid_assignment: boolean + If True then the class can assign sids where necessary. + constraints: boolean + If True, create SQL ForeignKey and Index constraints. + + """ + + self.allow_sid_assignment = allow_sid_assignment + if allow_sid_assignment: + ts = normalize_date(pd.Timestamp('now', tz='UTC')) + # Store as seconds since UNIX Epoch for compatibility + # with SQL. + self.end_date_to_assign = (ts.value // 10 ** 9) + + # Store a nested-dict of all metadata for + # reference when building Assets + self.metadata_cache = {} + + # Create SQL tables + self.init_db(db_conn, constraints) + + # Get the data to add to SQL + equities, futures, exchanges, root_symbols = self.load_data() + + # Write to the SQL tables + self._write_exchanges(exchanges, db_conn) + self._write_root_symbols(root_symbols, db_conn) + self._write_futures(futures, db_conn) + self._write_equities(equities, db_conn) + + def _write_exchanges(self, exchanges, db_conn): + + exchanges.to_sql('futures_exchanges', db_conn, if_exists='replace', + index=True, index_label='exchange_id') + + def _write_root_symbols(self, root_symbols, db_conn): + + root_symbols.to_sql('futures_root_symbols', db_conn, + if_exists='replace', index=True, + index_label='root_symbol_id') + + def _write_futures(self, futures, db_conn): + + futures.to_sql('futures_contracts', db_conn, if_exists='append', + index=True, index_label='sid') + + def _write_equities(self, equities, db_conn): + + equities.to_sql('equities', db_conn, if_exists='append', + index=True, index_label='sid') + + def init_db(self, + db_conn, + constraints=False): + """Connect to database and create tables. + + Parameters + ---------- + db_conn: sqlite3.Connection + A connection to a SQLite database. + constraints: boolean + If True, create SQL ForeignKey and Index constraints. + """ + + c = db_conn.cursor() + + c.execute(""" + CREATE TABLE IF NOT EXISTS equities ( + sid INTEGER NOT NULL, + symbol TEXT, + asset_name TEXT, + start_date INTEGER DEFAULT 0, + end_date INTEGER, + first_traded INTEGER, + exchange TEXT, + fuzzy TEXT + )""") + + c.execute(""" + CREATE TABLE IF NOT EXISTS futures_exchanges ( + exchange_id INTEGER NOT NULL, + exchange TEXT, + timezone TEXT + )""") + + c.execute(""" + CREATE TABLE IF NOT EXISTS futures_root_symbols ( + root_symbol_id INTEGER NOT NULL, + root_symbol TEXT, + sector TEXT, + description TEXT, + exchange_id INTEGER{fk} + )""".format(fk=", FOREIGN KEY(exchange_id) REFERENCES " + "futures_exchanges(exchange_id)" + if constraints else "")) + + c.execute(""" + CREATE TABLE IF NOT EXISTS futures_contracts ( + sid INTEGER NOT NULL, + symbol TEXT, + root_symbol_id INTEGER, + root_symbol TEXT, + asset_name TEXT, + start_date INTEGER DEFAULT 0, + end_date INTEGER, + first_traded INTEGER, + exchange_id INTEGER, + exchange TEXT, + notice_date INTEGER, + expiration_date INTEGER, + contract_multiplier REAL{fk} + )""".format(fk=", FOREIGN KEY(exchange_id) REFERENCES " + "futures_exchanges(exchange_id), " + "FOREIGN KEY(root_symbol_id) REFERENCES " + "futures_root_symbols(root_symbol_id)" + if constraints else "")) + + c.execute(""" + CREATE TABLE IF NOT EXISTS asset_router + (sid integer, + asset_type text + )""") + + if constraints: + + c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_equities_sid ' + 'ON equities(sid)') + c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_equities_symbol ' + 'ON equities(symbol)') + c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_equities_fuzzy ' + 'ON equities(fuzzy)') + c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_futures_exchanges_en ' # noqa + 'ON futures_exchanges(exchange_id)') + c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_futures_contracts_sid ' # noqa + 'ON futures_contracts(sid)') + c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_futures_root_symbols_id ' # noqa + 'ON futures_root_symbols(root_symbol_id)') + c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_asset_router_sid ' + 'ON asset_router(sid)') + + # Note: Also need a max_date table. + + db_conn.commit() + + @abstractmethod + def load_data(self): + """ + Subclasses should implement this method to return data in a standard + format: a pandas.DataFrame for each of the following tables: + equities, futures, exchanges, root_symbols + """ + + raise NotImplementedError('load_data') + + +class AssetDBWriterFromDictionary(AssetDBWriter): + """ + Class used to write dictionary data to SQLite database. + + Expects a dictionary to be passed to load_data + with the following format: + + {id_0: {start_date : ...}, id_1: {start_data: ...}, ...} + """ + + def __init__(self, equities={}, futures={}, exchanges={}, root_symbols={}): + + self._equities = equities + self._futures = futures + self._exchanges = exchanges + self._root_symbols = root_symbols + + def load_data(self): + """ + Convert our nested dictionaries to pandas DataFrames. + """ + equities_data = pd.DataFrame.from_dict(self._equities, orient='index') + + futures_data = pd.DataFrame.from_dict(self._futures, orient='index') + + exchange_data = pd.DataFrame.from_dict(self._exchanges, orient='index') + + root_symbol_data = pd.DataFrame.from_dict(self._root_symbols, + orient='index') + + # Assume the keys are the exchange_ids + exchange_cols = ['exchange', 'timezone'] + exchanges = pd.DataFrame(columns=exchange_cols) + + # Assume the keys are the root_symbol_ids + root_symbols_cols = ['root_symbol', 'sector', + 'description', 'exchange_id'] + root_symbols = pd.DataFrame(columns=root_symbols_cols) + + # Assume the keys are the sids + futures_cols = ['symbol', 'root_symbol', 'asset_name', + 'start_date', 'end_date', 'first_traded', 'exchange', + 'notice_date', 'expiration_date', + 'contract_multiplier'] + futures = pd.DataFrame(columns=futures_cols) + + # Assume the keys are the sids + equities_cols = ['symbol', 'asset_name', 'start_date', + 'end_date', 'first_traded', 'exchange', 'fuzzy'] + equities = pd.DataFrame(columns=equities_cols) + + # Append any data the user has provided. + exchanges = exchanges.append(exchange_data, verify_integrity=True) + root_symbols = root_symbols.append(root_symbol_data, + verify_integrity=True) + futures = futures.append(futures_data, verify_integrity=True) + equities = equities.append(equities_data, verify_integrity=True) + + return equities, futures, exchanges, root_symbols + + +class AssetDBWriterLegacy(AssetDBWriter): + """ + Overwrites some of the functionality of AssetDBWriter. + Used for backward compatibility. Will be deprecated. + + Methods + ------- + write_all(db_conn, fuzzy_char=None, allow_sid_assignment=True, + constraints=False) + Write the data supplied at initialization to the database. + write_block(self, identifier, **kwargs) + Inserts the given metadata kwargs to the entry for the given + sid. Matching fields in the existing entry will be overwritten. + Will be deprecated in future versions of zipline. + init_db(db_conn, constraints=False) + Create the SQLite tables (called by write_all). + load_data(self, equities, futures, exchanges, root_symbols) + Returns data in standard format. + consume_identifiers(self, db_conn, fuzzy_char=None, + allow_sid_assignment=True, + constraints=False) + Consume the identifiers supplied at initialization. + Will be deprecated in future versions of zipline. + """ + + def __init__(self, data): + + self._data = data + self.metadata_cache = {} + + def write_all(self, + db_conn, + fuzzy_char=None, + allow_sid_assignment=True, + constraints=False): + """Top-level entry point for writing a new asset db. + + Parameters + ---------- + db_conn: sqlite3.Connection + A connection to our SQLite database. + fuzzy_char: string + A string to be used in fuzzy matching. + allow_sid_assignment: boolean + If True, allow the writer to assign sids where necessary. + constraints: boolean + If True, add SQL constraints to tables. + """ + + self.conn = db_conn + + self.fuzzy_char = fuzzy_char + self.allow_sid_assignment = allow_sid_assignment + + # This flag controls if the AssetDBWriter is allowed to generate its + # own sids. If False, metadata that does not contain a sid will raise + # an exception when building assets. + if allow_sid_assignment: + self.end_date_to_assign = normalize_date( + pd.Timestamp('now', tz='UTC')) + + # Create SQL tables. + self.init_db(self.conn, constraints) + + # Write to SQL tables. + for sid, metadata in self.load_data(self._data): + self.write_block(sid, **metadata) + + def write_block(self, identifier, **kwargs): + """ + Inserts the given metadata kwargs to the entry for the given + sid. Matching fields in the existing entry will be overwritten. + Will be deprecated in future versions of zipline. + """ + + 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 + if key not in ASSET_FIELDS: + continue + # Do not accept Nones + if value is None: + continue + # Do not accept empty strings + if value == '': + continue + # Do not accept NaNs from dataframes + if isinstance(value, float) and np.isnan(value): + continue + entry[key] = value + + # Check if the sid is declared + try: + entry['sid'] + except KeyError: + # If the sid is not a sid, assign one + if hasattr(identifier, '__int__'): + entry['sid'] = identifier.__int__() + else: + if self.allow_sid_assignment: + # Assign the sid the value of its insertion order. + # This assumes that we are assigning values to all assets. + entry['sid'] = len(self.metadata_cache) + 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': + try: + fuzzy = entry['symbol'].replace(self.fuzzy_char, '') \ + if self.fuzzy_char else None + except KeyError: + fuzzy = 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_contracts( + 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 + + self.conn.commit() + + def consume_identifiers(self, db_conn, fuzzy_char=None, + allow_sid_assignment=True, constraints=False): + """ + Consumes the given identifiers in to the metadata cache of this + AssetDBWriter, and adds to database. + Will be deprecated in future versions of zipline. + """ + + self.conn = db_conn + self.fuzzy_char = fuzzy_char + self.allow_sid_assignment = allow_sid_assignment + + # This flag controls if the AssetDBWriter is allowed to generate its + # own sids. If False, metadata that does not contain a sid will raiset + # an exception when building assets. + if allow_sid_assignment: + self.end_date_to_assign = normalize_date( + pd.Timestamp('now', tz='UTC')) + + # Create SQL tables + self.init_db(self.conn, constraints) + + for identifier in self._data: + # Handle case where full Assets are passed in + # For example, in the creation of a DataFrameSource, the source's + # 'sid' args may be full Assets + if isinstance(identifier, Asset): + sid = identifier.sid + metadata = identifier.to_dict() + metadata['asset_type'] = identifier.__class__.__name__ + self.write_block(sid, **metadata) + else: + self.write_block(identifier) + + +class NullAssetDBWriterLegacy(AssetDBWriterLegacy): + """ + An implementation of AssetDBWriterLegacy for use + when no data is initially specified. + """ + + def load_data(self, __): + for i in iter(()): + yield + + +class AssetDBWriterLegacyFromList(AssetDBWriterLegacy): + """ + Returns a generator yielding entries from sid_list. + """ + def load_data(self, sid_list): + + for i in sid_list: + yield i + + +class AssetDBWriterLegacyFromDictionary(AssetDBWriterLegacy): + """ An implementation of AssetDBWriter for use + with dictionaries. + + Expects a dictionary to be passed to load_data + with the following format: + + {id_0: {start_date : ...}, id_1: {start_data: ...}, ...} + """ + + def load_data(self, dict): + """ + Returns a generator yielding pairs of (identifier, metadata) + """ + for identifier, metadata in dict.items(): + yield identifier, metadata + + +class AssetDBWriterLegacyFromDataFrame(AssetDBWriterLegacy): + """ An implementation of AssetDBWriter for use + with pandas DataFrames. + + Expects dataframe to be passed to load_data + to have the following structure: + * column names must be the metadata fields + * index must be the different asset identifiers + * array contents should be the metadata value + """ + + def load_data(self, dataframe): + """ + Returns a generator yielding pairs of (identifier, metadata) + """ + for identifier, row in dataframe.iterrows(): + yield identifier, row.to_dict() + + +class AssetDBWriterLegacyFromReadable(AssetDBWriterLegacy): + """ An implementation of AssetDBWriter for use + with objects with a 'read' property. + + The object's read method must return rows + containing at least one of 'sid' or 'symbol' along + with the other metadata fields. + """ + + def load_data(self, readable): + """ + Returns a generator yielding pairs of (identifier, metadata) + """ + for row in readable.read(): + id_metadata = {} + for field in ASSET_FIELDS: + try: + row_value = row[field] + # Avoid passing placeholder strings + if row_value and (row_value != 'None'): + id_metadata[field] = row[field] + except KeyError: + continue + except IndexError: + continue + if 'sid' in id_metadata: + identifier = id_metadata['sid'] + del id_metadata['sid'] + elif 'symbol' in id_metadata: + identifier = id_metadata['symbol'] + del id_metadata['symbol'] + else: + raise ConsumeAssetMetaDataError(obj=row) + yield identifier, id_metadata diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index 612b847c..8a95f556 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -14,11 +14,11 @@ from abc import ABCMeta from numbers import Integral -import numpy as np import sqlite3 from sqlite3 import Row import warnings +import numpy as np from logbook import Logger import pandas as pd from pandas.tseries.tools import normalize_date @@ -26,112 +26,87 @@ from six import with_metaclass, string_types from zipline.errors import ( ConsumeAssetMetaDataError, - InvalidAssetType, MultipleSymbolsFound, RootSymbolNotFound, - SidAssignmentError, SidNotFound, SymbolNotFound, MapAssetIdentifierIndexError, ) -from zipline.assets._assets import ( +from zipline.assets import ( Asset, Equity, Future ) +from zipline.assets import ( + NullAssetDBWriterLegacy, + AssetDBWriterLegacyFromList, + AssetDBWriterLegacyFromDictionary, + AssetDBWriterLegacyFromDataFrame, + AssetDBWriterLegacyFromReadable +) +from zipline.assets.asset_writer import ( + FUTURE_TABLE_FIELDS, + EQUITY_TABLE_FIELDS +) log = Logger('assets.py') -# Expected fields for an Asset's metadata -ASSET_FIELDS = [ - 'sid', - 'asset_type', - 'symbol', - 'root_symbol', - 'asset_name', - 'start_date', - 'end_date', - 'first_traded', - 'exchange', - 'notice_date', - 'expiration_date', - 'contract_multiplier', - # The following fields are for compatibility with other systems - 'file_name', # Used as symbol - 'company_name', # Used as asset_name - 'start_date_nano', # Used as start_date - 'end_date_nano', # Used as end_date -] - - -# 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( +FUTURE_BY_SID_QUERY = 'select {0} from futures_contracts where sid=?'.format( ", ".join(FUTURE_TABLE_FIELDS)) EQUITY_BY_SID_QUERY = 'select {0} from equities where sid=?'.format( ", ".join(EQUITY_TABLE_FIELDS)) +def create_relevant_writer(metadata): + """ Create an instance of AssetDBWriter relevant to + processing metadata. + Will be deprecated in future versions of zipline. + """ + + if isinstance(metadata, dict): + return AssetDBWriterLegacyFromDictionary(metadata) + elif isinstance(metadata, pd.DataFrame): + return AssetDBWriterLegacyFromDataFrame(metadata) + elif isinstance(metadata, list): + return AssetDBWriterLegacyFromList(metadata) + elif hasattr(metadata, 'read'): + return AssetDBWriterLegacyFromReadable(metadata) + elif metadata is None: + return NullAssetDBWriterLegacy(metadata) + else: + raise ConsumeAssetMetaDataError(obj=metadata) + + class AssetFinder(object): - def __init__(self, - metadata=None, - allow_sid_assignment=True, - fuzzy_char=None, - db_path=':memory:', - create_table=True): + def __init__(self, metadata=None, allow_sid_assignment=True, + fuzzy_char=None, db_path=':memory:', create_table=True, + asset_writer=None): 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 - if allow_sid_assignment: - self.end_date_to_assign = normalize_date( - pd.Timestamp('now', tz='UTC')) - self.conn = sqlite3.connect(db_path) - self.conn.text_factory = str - self.cursor = self.conn.cursor() - # The AssetFinder also holds a nested-dict of all metadata for - # reference when building Assets - self.metadata_cache = {} + # AssetFinder can optionally accept an instance of + # the AssetDBWriter class. If no writer is supplied, + # we create a relevant writer based on the supplied metadata. + # Note that this strucutre is for backward compatibility. + # Ultimately AssetDBWriter and AssetFinder will be completely + # separate, and AssetFinder will not instantiate AssetDBWriter. + if asset_writer is None: + _asset_writer = create_relevant_writer(metadata) + else: + _asset_writer = asset_writer - # 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. + # Create tables and read in metadata. if create_table: - self.create_db_tables() + _asset_writer.init_db(self.conn) if metadata is not None: - self.consume_metadata(metadata) + _asset_writer.write_all(self.conn, + self.fuzzy_char, + self.allow_sid_assignment) # 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. @@ -151,54 +126,24 @@ class AssetFinder(object): # Populated on first call to `lifetimes`. self._asset_lifetimes = None - 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 clear_metadata(self): + """ + Used for testing. + Will be deprecated in future versions of zipline. + """ + # Close the database connection + self.conn.close() + # Create new database connection in memory. + self.conn = sqlite3.connect(':memory:') + # Initialize the database tables using the same connection + # as used by the AssetFinder. + _asset_writer = NullAssetDBWriterLegacy({}) + _asset_writer.init_db(self.conn) def asset_type_by_sid(self, sid): + """ + Retrieve the asset type of a given sid. + """ try: return self._asset_type_cache[sid] except KeyError: @@ -207,7 +152,7 @@ class AssetFinder(object): 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' + query = 'SELECT asset_type FROM asset_router WHERE sid=:sid' c.execute(query, t) data = c.fetchone() if data is None: @@ -219,6 +164,9 @@ class AssetFinder(object): return asset_type def retrieve_asset(self, sid, default_none=False): + """ + Retrieve the Asset object of a given sid. + """ if isinstance(sid, Asset): return sid @@ -246,6 +194,9 @@ class AssetFinder(object): return [self.retrieve_asset(sid) for sid in sids] def _retrieve_equity(self, sid): + """ + Retrieve the Equity object of a given sid. + """ try: return self._equity_cache[sid] except KeyError: @@ -275,6 +226,9 @@ class AssetFinder(object): return equity def _retrieve_futures_contract(self, sid): + """ + Retrieve the Future object of a given sid. + """ try: return self._future_cache[sid] except KeyError: @@ -329,10 +283,10 @@ class AssetFinder(object): 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>=?") + query = ("SELECT sid FROM equities " + "WHERE symbol=? " + "AND start_date<=? " + "AND end_date>=?") c.execute(query, t) candidates = c.fetchall() @@ -343,11 +297,11 @@ class AssetFinder(object): # 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") + query = ("SELECT sid FROM equities " + "WHERE symbol=? " + "AND start_date<=? " + "ORDER BY end_date DESC " + "LIMIT 1") c.execute(query, t) data = c.fetchone() @@ -358,11 +312,11 @@ class AssetFinder(object): # 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") + 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() @@ -373,7 +327,7 @@ class AssetFinder(object): else: t = (symbol,) - query = ("select sid from equities where symbol=?") + query = ("SELECT sid FROM equities WHERE symbol=?") c.execute(query, t) data = c.fetchall() @@ -411,10 +365,10 @@ class AssetFinder(object): 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>=?") + query = ("SELECT sid FROM equities " + "WHERE fuzzy=? " + + "AND start_date<=? " + + "AND end_date>=?") c.execute(query, t) candidates = c.fetchall() @@ -426,11 +380,11 @@ class AssetFinder(object): # 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") + 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: @@ -469,7 +423,6 @@ class AssetFinder(object): root symbol. """ c = self.conn.cursor() - if as_of_date is pd.NaT: # If the as_of_date is NaT, get all contracts for this # root symbol. @@ -501,7 +454,8 @@ class AssetFinder(object): if not sids: # Check if root symbol exists. c.execute(""" - select count(sid) from futures where root_symbol=:root_symbol + SELECT COUNT(sid) FROM futures_contracts + WHERE root_symbol=:root_symbol """, t) count = c.fetchone()[0] if count == 0: @@ -514,7 +468,7 @@ class AssetFinder(object): @property def sids(self): c = self.conn.cursor() - query = 'select sid from asset_router' + query = 'SELECT sid FROM asset_router' c.execute(query) return [r[0] for r in c.fetchall()] @@ -615,14 +569,14 @@ class AssetFinder(object): input index. Parameters - __________ + ---------- index : Iterable An iterable containing ints, strings, or Assets as_of_date : pandas.Timestamp A date to be used to resolve any dual-mapped symbols Returns - _______ + ------- List A list of integer sids corresponding to the input index """ @@ -656,274 +610,49 @@ 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): - """ - 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 - """ - 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 - if key not in ASSET_FIELDS: - continue - # Do not accept Nones - if value is None: - continue - # Do not accept empty strings - if value == '': - continue - # Do not accept nans from dataframes - if isinstance(value, float) and np.isnan(value): - continue - entry[key] = value - - # Check if the sid is declared - try: - entry['sid'] - except KeyError: - # If the identifier is not a sid, assign one - if hasattr(identifier, '__int__'): - entry['sid'] = identifier.__int__() - else: - if self.allow_sid_assignment: - # Assign the sid the value of its insertion order. - # This assumes that we are assigning values to all assets. - entry['sid'] = len(self.metadata_cache) - 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': - try: - fuzzy = entry['symbol'].replace(self.fuzzy_char, '') \ - if self.fuzzy_char else None - except KeyError: - fuzzy = 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): """ - Consumes the given identifiers in to the metadata cache of this - AssetFinder. + Consumes the provided identifiers, passing them to + the asset writer to be added to the database. + Will be deprecated in future versions of zipline. + + Parameters + ---------- + identifiers + The data to be consumed. """ - for identifier in identifiers: - # Handle case where full Assets are passed in - # For example, in the creation of a DataFrameSource, the source's - # 'sid' args may be full Assets - if isinstance(identifier, Asset): - sid = identifier.sid - metadata = identifier.to_dict() - metadata['asset_type'] = identifier.__class__.__name__ - self.insert_metadata(identifier=sid, **metadata) - else: - self.insert_metadata(identifier) + _asset_writer = AssetDBWriterLegacyFromList(identifiers) + _asset_writer.consume_identifiers(self.conn, + self.fuzzy_char, + self.allow_sid_assignment) def consume_metadata(self, metadata): """ - Consumes the provided metadata in to the metadata cache. The - existing values in the cache will be overwritten when there - is a conflict. - :param metadata: The metadata to be consumed - """ - # Handle dicts - if isinstance(metadata, dict): - self._insert_metadata_dict(metadata) - # Handle DataFrames - elif isinstance(metadata, pd.DataFrame): - self._insert_metadata_dataframe(metadata) - # Handle readables - elif hasattr(metadata, 'read'): - self._insert_metadata_readable(metadata) - else: - raise ConsumeAssetMetaDataError(obj=metadata) + Consumes the provided metadata, passing it to + the asset writer to be added to the database. + Will be deprecated in future versions of zipline. - def clear_metadata(self): + Parameters + ---------- + metadata + The data to be consumed. """ - Used for testing. - """ - self.metadata_cache = {} - - self.conn = sqlite3.connect(':memory:') - self.create_db_tables() + _asset_writer = create_relevant_writer(metadata) + _asset_writer.write_all(self.conn, + fuzzy_char=self.fuzzy_char, + allow_sid_assignment=self.allow_sid_assignment) 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.conn.commit() - - def _insert_metadata_dict(self, dict): - for identifier, entry in dict.items(): - self._insert_metadata(identifier, **entry) - self.conn.commit() - - def _insert_metadata_readable(self, readable): - for row in readable.read(): - # Parse out the row of the readable object - metadata_dict = {} - for field in ASSET_FIELDS: - try: - row_value = row[field] - # Avoid passing placeholders - if row_value and (row_value != 'None'): - metadata_dict[field] = row[field] - except KeyError: - continue - except IndexError: - continue - # Locate the identifier, fail if not found - if 'sid' in metadata_dict: - identifier = metadata_dict['sid'] - elif 'symbol' in metadata_dict: - identifier = metadata_dict['symbol'] - else: - raise ConsumeAssetMetaDataError(obj=row) - self._insert_metadata(identifier, **metadata_dict) - self.conn.commit() + """ + Insert information for a single identifier. + Will be deprecated in future versions of zipline. + """ + metadata = {} + metadata[identifier] = kwargs + _asset_writer = create_relevant_writer(metadata) + _asset_writer.write_all(self.conn, + fuzzy_char=self.fuzzy_char, + allow_sid_assignment=self.allow_sid_assignment) def _compute_asset_lifetimes(self): """