wassname/catalyst
1
0
Fork 0
mirror of https://github.com/wassname/catalyst.git synced 2026-09-12 12:12:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

ENH: Create AssetDBWriter class

Browse Source 
The AssetDBWriter class and its subclasses will
ultimately be responsible for creating the SQLite
database tables and writing data to these tables.

In the longer term AssetDBWriter and AssetFinder will
be decoupled, sharing only an SQLite connection.
However, for backward compatibility reasons this has
not yet been fully implemented.

Modify tests since AssetFinder no longer has a
metadata_cache attribute.
This commit is contained in:
Stewart Douglas
2015-09-10 11:53:22 -04:00
committed by jfkirk
parent 45adc57267
commit 746f70a133
5 changed files with 870 additions and 414 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_assets.py
+4 -3
View File
@@ -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
zipline/assets/__init__.py
+14
View File
@@ -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',
zipline/assets/_assets.pyx
+1
View File
@@ -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,
zipline/assets/asset_writer.py
+711
View File
@@ -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
zipline/assets/assets.py
+140 -411
View File
@@ -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):
"""