mirror of
https://github.com/wassname/catalyst.git
synced 2026-08-15 12:15:22 +08:00
ENH: sqlalchemy
This commit is contained in:
llllllllll
@@ -45,3 +45,5 @@ bcolz==0.10.0
|
||||
|
||||
# Command line interface helper
|
||||
click==4.0.0
|
||||
|
||||
toolz==0.7.2
|
||||
|
||||
+148
-210
@@ -3,10 +3,11 @@ from abc import (
|
||||
abstractmethod,
|
||||
)
|
||||
|
||||
import pandas as pd
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from pandas.tseries.tools import normalize_date
|
||||
from six import with_metaclass, string_types
|
||||
import sqlalchemy as sa
|
||||
|
||||
from zipline.errors import (
|
||||
ConsumeAssetMetaDataError,
|
||||
@@ -17,7 +18,7 @@ from zipline.assets import (
|
||||
Asset, Equity, Future
|
||||
)
|
||||
|
||||
ASSET_FIELDS = [
|
||||
ASSET_FIELDS = frozenset({
|
||||
'sid',
|
||||
'asset_type',
|
||||
'symbol',
|
||||
@@ -34,43 +35,44 @@ ASSET_FIELDS = [
|
||||
'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
|
||||
]
|
||||
'end_date_nano', # Used as end_date
|
||||
})
|
||||
|
||||
# Expected fields for an Asset's metadata
|
||||
ASSET_TABLE_FIELDS = [
|
||||
ASSET_TABLE_FIELDS = frozenset({
|
||||
'sid',
|
||||
'symbol',
|
||||
'asset_name',
|
||||
'start_date',
|
||||
'end_date',
|
||||
'first_traded',
|
||||
'exchange'
|
||||
]
|
||||
'exchange',
|
||||
})
|
||||
|
||||
|
||||
# Expected fields for an Asset's metadata
|
||||
FUTURE_TABLE_FIELDS = ASSET_TABLE_FIELDS + [
|
||||
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_TABLE_FIELDS = frozenset({
|
||||
'exchange_id',
|
||||
'exchange',
|
||||
'timezone'
|
||||
]
|
||||
})
|
||||
|
||||
ROOT_SYMBOL_TABLE_FIELDS = [
|
||||
ROOT_SYMBOL_TABLE_FIELDS = ({
|
||||
'root_symbol_id',
|
||||
'root_symbol',
|
||||
'sector',
|
||||
'description',
|
||||
'exchange_id'
|
||||
]
|
||||
})
|
||||
|
||||
|
||||
class AssetDBWriter(with_metaclass(ABCMeta)):
|
||||
@@ -81,237 +83,173 @@ class AssetDBWriter(with_metaclass(ABCMeta)):
|
||||
|
||||
Methods
|
||||
-------
|
||||
write_all(db_conn, fuzzy_char=None, allow_sid_assignment=True,
|
||||
write_all(engine, fuzzy_char=None, allow_sid_assignment=True,
|
||||
constraints=False)
|
||||
Write the data supplied at initialization to the database.
|
||||
init_db(db_conn, constraints=False)
|
||||
init_db(engine, constraints=False)
|
||||
Create the SQLite tables (called by write_all).
|
||||
load_data()
|
||||
Returns data in standard format.
|
||||
|
||||
"""
|
||||
def __init__(self):
|
||||
self.sql_metadata = None
|
||||
|
||||
def write_all(self, db_conn, fuzzy_char=None, allow_sid_assignment=True,
|
||||
def write_all(self,
|
||||
engine,
|
||||
fuzzy_char=None,
|
||||
constraints=False):
|
||||
""" Write pre-supplied data to SQLite.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
db_conn: sqlite3.Connection
|
||||
A connection to a SQLite database.
|
||||
fuzzy_char: string
|
||||
engine : Engine
|
||||
An engine to a SQL database.
|
||||
fuzzy_char : str, optional
|
||||
A string for use in fuzzy matching.
|
||||
allow_sid_assignment: boolean
|
||||
allow_sid_assignment: bool, optional
|
||||
If True then the class can assign sids where necessary.
|
||||
constraints: boolean
|
||||
constraints : bool, optional
|
||||
If True, create SQL ForeignKey and Index constraints.
|
||||
|
||||
"""
|
||||
self.fuzzy_char = fuzzy_char
|
||||
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)
|
||||
|
||||
self.init_db(engine, constraints)
|
||||
# Get the data to add to SQL
|
||||
equities, futures, exchanges, root_symbols = self.load_data()
|
||||
with engine.begin() as txn:
|
||||
self._write_exchanges(exchanges, txn)
|
||||
self._write_root_symbols(root_symbols, txn)
|
||||
self._write_futures(futures, txn)
|
||||
self._write_equities(equities, txn)
|
||||
|
||||
c = db_conn.cursor()
|
||||
# Write to the SQL tables, using the raw SQL driver instead
|
||||
# of the pandas.DataFrame.to_sql method, as the former
|
||||
# allows us to create an SQL transaction
|
||||
c.execute('BEGIN')
|
||||
# Everything between here and the db_conn.commit()
|
||||
# will be part of of the same SQL transaction.
|
||||
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)
|
||||
db_conn.commit()
|
||||
def _write_exchanges(self, exchanges, bind=None):
|
||||
self.futures_exchanges.insert().values(
|
||||
exchanges.to_records(),
|
||||
).execute(bind=bind)
|
||||
|
||||
def _write_exchanges(self, exchanges, db_conn):
|
||||
def _write_root_symbols(self, root_symbols, bind=None):
|
||||
self.futures_root_symbols.insert().values(
|
||||
root_symbols.to_records(),
|
||||
).execute(bind=bind)
|
||||
|
||||
data = [tuple(x) for x in exchanges.to_records()]
|
||||
|
||||
c = db_conn.cursor()
|
||||
# The OR IGNORE syntax means we do not insert data
|
||||
# which would violate an SQL constraint.
|
||||
c.executemany("""
|
||||
INSERT OR IGNORE INTO futures_exchanges
|
||||
('exchange_id', 'exchange', 'timezone')
|
||||
VALUES (?, ?, ?)
|
||||
""", data)
|
||||
|
||||
def _write_root_symbols(self, root_symbols, db_conn):
|
||||
|
||||
data = [tuple(x) for x in root_symbols.to_records()]
|
||||
|
||||
c = db_conn.cursor()
|
||||
c.executemany("""
|
||||
INSERT OR IGNORE INTO futures_root_symbols
|
||||
('root_symbol_id', 'root_symbol', 'sector',
|
||||
'description', 'exchange_id')
|
||||
VALUES(?, ?, ?, ?, ?)
|
||||
""", data)
|
||||
|
||||
def _write_futures(self, futures, db_conn):
|
||||
|
||||
# Retrieve the data to add to the futures_contracts table
|
||||
data = [tuple(x) for x in futures.to_records()]
|
||||
|
||||
# Retrieve the data to add to the asset_router table
|
||||
sids = [x[0] for x in data]
|
||||
futs = ['future'] * len(sids)
|
||||
to_insert = zip(sids, futs)
|
||||
|
||||
c = db_conn.cursor()
|
||||
c.executemany("""
|
||||
INSERT OR IGNORE INTO futures_contracts
|
||||
('sid', 'symbol', 'root_symbol', 'asset_name',
|
||||
'start_date', 'end_date', 'first_traded', 'exchange',
|
||||
'notice_date', 'expiration_date', 'contract_multiplier')
|
||||
VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", data)
|
||||
|
||||
c.executemany("""
|
||||
INSERT INTO asset_router
|
||||
('sid', 'asset_type')
|
||||
VALUES(?,?)
|
||||
""", to_insert)
|
||||
|
||||
def _write_equities(self, equities, db_conn):
|
||||
def _write_futures(self, futures, bind=None):
|
||||
recs = futures.to_records()
|
||||
self.futures_contracts.insert().values(recs).execute(bind=bind)
|
||||
self.asset_router.insert().values([
|
||||
(rec['sid'], 'future') for rec in recs
|
||||
]).execute(bind=bind)
|
||||
|
||||
def _write_equities(self, equities, fuzzy_char, bind=None):
|
||||
# Apply fuzzy matching.
|
||||
if self.fuzzy_char:
|
||||
equities['fuzzy'] = equities['symbol'].str.\
|
||||
replace(self.fuzzy_char, '')
|
||||
if fuzzy_char:
|
||||
equities['fuzzy'] = equities['symbol'].str.replace(fuzzy_char, '')
|
||||
|
||||
# Retrieve the data to add to the equitites table
|
||||
data = [tuple(x) for x in equities.to_records()]
|
||||
recs = equities.to_records()
|
||||
self.equities.insert().values(recs).execute(bind=bind)
|
||||
self.asset_router.insert().values([
|
||||
(rec['sid'], 'equity') for rec in recs
|
||||
]).execute(bind=bind)
|
||||
|
||||
# Retrieve the data to add to the asset_router table
|
||||
sids = [x[0] for x in data]
|
||||
eqs = ['equity'] * len(sids)
|
||||
to_insert = zip(sids, eqs)
|
||||
|
||||
c = db_conn.cursor()
|
||||
c.executemany("""
|
||||
INSERT OR IGNORE INTO equities
|
||||
('sid', 'symbol', 'asset_name', 'start_date',
|
||||
'end_date', 'first_traded', 'exchange', 'fuzzy')
|
||||
VALUES(?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", data)
|
||||
|
||||
c.executemany("""
|
||||
INSERT INTO asset_router
|
||||
('sid', 'asset_type')
|
||||
VALUES(?,?)
|
||||
""", to_insert)
|
||||
|
||||
def init_db(self,
|
||||
db_conn,
|
||||
constraints=False):
|
||||
def init_db(self, engine, constraints=False):
|
||||
"""Connect to database and create tables.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
db_conn: sqlite3.Connection
|
||||
A connection to a SQLite database.
|
||||
constraints: boolean
|
||||
engine : Engine
|
||||
An engine to a SQL database.
|
||||
constraints : bool
|
||||
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
|
||||
)""")
|
||||
|
||||
# Note: Would be optimal to use INTEGER PRIMARY KEY here, but SQLite
|
||||
# tables cannot be modified after creation. Using CREATE UNIQUE INDEX
|
||||
# only marginally less performant.
|
||||
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()
|
||||
self.sql_metadata = metadata = sa.MetaData(bind=engine)
|
||||
self.equities = sa.Table(
|
||||
'equities',
|
||||
metadata,
|
||||
sa.Column(
|
||||
'sid',
|
||||
sa.Integer,
|
||||
nullable=False, # in case constraints is False
|
||||
primary_key=constraints,
|
||||
),
|
||||
sa.Column('symbol', sa.Text),
|
||||
sa.Column('asset_name', sa.Text),
|
||||
sa.Column('start_date', sa.Integer, default=0),
|
||||
sa.Column('end_date', sa.Integer),
|
||||
sa.Column('first_traded', sa.Integer),
|
||||
sa.Column('exchange', sa.Text),
|
||||
sa.Column('fuzzy', sa.Text),
|
||||
)
|
||||
self.futures_exchanges = sa.Table(
|
||||
'futures_exchanges',
|
||||
metadata,
|
||||
sa.Column(
|
||||
'exchange_id',
|
||||
sa.Integer,
|
||||
nullable=False, # in case constraints is False
|
||||
primary_key=constraints,
|
||||
),
|
||||
sa.Column('exchange', sa.Text),
|
||||
sa.Column('timezone', sa.Text),
|
||||
)
|
||||
self.futures_root_symbols = sa.Table(
|
||||
'futures_root_symbols',
|
||||
metadata,
|
||||
sa.Column(
|
||||
'root_symbol_id',
|
||||
sa.Integer,
|
||||
nullable=False, # in case constraints is False
|
||||
primary_key=constraints,
|
||||
),
|
||||
sa.Column('root_symbol', sa.Text),
|
||||
sa.Column('sector', sa.Text),
|
||||
sa.Column('description', sa.Text),
|
||||
sa.Column(
|
||||
'exchange_id',
|
||||
sa.Integer,
|
||||
*((sa.ForeignKey(self.futures_exchanges.c.exchange_id),)
|
||||
if constraints else ())
|
||||
),
|
||||
)
|
||||
self.futures_contracts = sa.Table(
|
||||
'futures_contracts',
|
||||
metadata,
|
||||
sa.Column(
|
||||
'sid',
|
||||
sa.Integer,
|
||||
nullable=False, # in case constraints is False
|
||||
primary_key=constraints,
|
||||
),
|
||||
sa.Column('symbol', sa.Text),
|
||||
sa.Column(
|
||||
'root_symbol_id',
|
||||
sa.Integer,
|
||||
*((sa.ForeignKey(self.futures_root_symbols.c.root_symbol_id),)
|
||||
if constraints else ())
|
||||
),
|
||||
sa.Column('root_symbol', sa.Text),
|
||||
sa.Column('asset_name', sa.Text),
|
||||
sa.Column('start_date', sa.Integer, default=0),
|
||||
sa.Column('end_date', sa.Integer),
|
||||
sa.Column('first_traded', sa.Integer),
|
||||
sa.Column(
|
||||
'exchange_id',
|
||||
sa.Integer,
|
||||
*((sa.ForeignKey(self.futures_exchanges.c.exchange_id),)
|
||||
if constraints else ())
|
||||
),
|
||||
sa.column('exchange', sa.Text),
|
||||
sa.Column('notice_date', sa.Integer),
|
||||
sa.Column('expiration_date', sa.Integer),
|
||||
sa.Column('contract_multiplier', sa.Float),
|
||||
)
|
||||
self.asset_router = sa.Table(
|
||||
'asset_router',
|
||||
metadata,
|
||||
sa.Column('sid', sa.Integer, primary_key=constraints),
|
||||
sa.Column('asset_type', sa.Text),
|
||||
)
|
||||
metadata.create_all(checkfirst=True)
|
||||
return metadata
|
||||
|
||||
@staticmethod
|
||||
def dict_subset(dict_, subset):
|
||||
|
||||
+198
-177
@@ -13,15 +13,18 @@
|
||||
# limitations under the License.
|
||||
|
||||
from abc import ABCMeta
|
||||
from functools import partial, itemgetter
|
||||
from numbers import Integral
|
||||
from sqlite3 import Row
|
||||
from operator import getitem
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
from logbook import Logger
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from pandas.tseries.tools import normalize_date
|
||||
from six import with_metaclass, string_types
|
||||
import sqlalchemy as sa
|
||||
from toolz import compose
|
||||
|
||||
from zipline.errors import (
|
||||
MultipleSymbolsFound,
|
||||
@@ -31,34 +34,74 @@ from zipline.errors import (
|
||||
MapAssetIdentifierIndexError,
|
||||
)
|
||||
from zipline.assets import (
|
||||
Asset, Equity, Future
|
||||
Asset, Equity, Future,
|
||||
)
|
||||
from zipline.assets.asset_writer import (
|
||||
FUTURE_TABLE_FIELDS,
|
||||
EQUITY_TABLE_FIELDS
|
||||
EQUITY_TABLE_FIELDS,
|
||||
)
|
||||
|
||||
log = Logger('assets.py')
|
||||
|
||||
# Create the query once from the fields, so that the join is not done
|
||||
# repeatedly.
|
||||
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))
|
||||
|
||||
|
||||
class AssetFinder(object):
|
||||
|
||||
def __init__(self, conn, allow_sid_assignment=True,
|
||||
fuzzy_char=None):
|
||||
def __init__(self, engine, allow_sid_assignment=True, fuzzy_char=None):
|
||||
|
||||
self.fuzzy_char = fuzzy_char
|
||||
self.allow_sid_assignment = allow_sid_assignment
|
||||
|
||||
self.conn = conn
|
||||
self.engine = engine
|
||||
metadata = sa.Metadata(bind=engine)
|
||||
self.equities = equities = sa.Table(
|
||||
'equities',
|
||||
metadata,
|
||||
autoload_with=engine,
|
||||
)
|
||||
self.futures_exchanges = sa.Table(
|
||||
'futures_exchanges',
|
||||
metadata,
|
||||
autoload_with=engine,
|
||||
)
|
||||
self.futures_root_symbols = sa.Table(
|
||||
'futures_root_symbols',
|
||||
metadata,
|
||||
autoload_with=engine,
|
||||
)
|
||||
self.futures_contracts = futures_contracts = sa.Table(
|
||||
'futures_contracts',
|
||||
metadata,
|
||||
autoload_with=engine,
|
||||
)
|
||||
self.asset_router = sa.Table(
|
||||
'asset_router',
|
||||
metadata,
|
||||
autoload_with=engine,
|
||||
)
|
||||
|
||||
# Create the equity and future queries once.
|
||||
_equity_sid = equities.c.sid
|
||||
_equity_by_sid = sa.select(
|
||||
tuple(map(partial(getitem, equities.c), EQUITY_TABLE_FIELDS)),
|
||||
)
|
||||
|
||||
def select_equity_by_sid(sid):
|
||||
return _equity_by_sid.where(_equity_sid == int(sid))
|
||||
|
||||
self.select_equity_by_sid = select_equity_by_sid
|
||||
|
||||
_future_sid = futures_contracts.c.sid
|
||||
_future_by_sid = sa.select(
|
||||
tuple(map(
|
||||
partial(getitem, futures_contracts.c),
|
||||
FUTURE_TABLE_FIELDS,
|
||||
)),
|
||||
)
|
||||
|
||||
def select_future_by_sid(sid):
|
||||
return _future_by_sid.where(_future_sid == int(sid))
|
||||
|
||||
self.select_future_by_sid = select_future_by_sid
|
||||
# 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.
|
||||
#
|
||||
@@ -86,18 +129,12 @@ class AssetFinder(object):
|
||||
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
|
||||
asset_type = sa.select((self.asset_router.c.asset_type,)).where(
|
||||
self.asset_router.c.sid == int(sid),
|
||||
).scalar().execute()
|
||||
|
||||
if asset_type is not None:
|
||||
self._asset_type_cache[sid] = asset_type
|
||||
return asset_type
|
||||
|
||||
def retrieve_asset(self, sid, default_none=False):
|
||||
@@ -139,11 +176,7 @@ class AssetFinder(object):
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
c = self.conn.cursor()
|
||||
c.row_factory = Row
|
||||
t = (int(sid),)
|
||||
c.execute(EQUITY_BY_SID_QUERY, t)
|
||||
data = dict(c.fetchone())
|
||||
data = self.select_equity_by_sid(sid).execute().fetchone()
|
||||
if data:
|
||||
if data['start_date']:
|
||||
data['start_date'] = pd.Timestamp(data['start_date'], tz='UTC')
|
||||
@@ -171,11 +204,7 @@ class AssetFinder(object):
|
||||
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())
|
||||
data = self.select_future_by_sid(sid).execute().fetchone()
|
||||
if data:
|
||||
if data['start_date']:
|
||||
data['start_date'] = pd.Timestamp(data['start_date'], tz='UTC')
|
||||
@@ -209,79 +238,69 @@ class AssetFinder(object):
|
||||
If multiple Assets are found and as_of_date is not set,
|
||||
raises MultipleSymbolsFound.
|
||||
|
||||
If no Asset was active at as_of_date, and allow_expired is False
|
||||
raises SymbolNotFound.
|
||||
If no Asset was active at as_of_date raises SymbolNotFound.
|
||||
"""
|
||||
if as_of_date is not None:
|
||||
as_of_date = pd.Timestamp(normalize_date(as_of_date))
|
||||
|
||||
c = self.conn.cursor()
|
||||
|
||||
equities_cols = self.equities.c
|
||||
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()
|
||||
ad_value = as_of_date.value
|
||||
|
||||
# If one SID exists for symbol, return that symbol
|
||||
candidates = sa.select((equities_cols.sid,)).where(
|
||||
(equities_cols.symbol == symbol) &
|
||||
(equities_cols.start_date <= ad_value) &
|
||||
(equities_cols.end_date >= ad_value),
|
||||
).execute().fetchall()
|
||||
if len(candidates) == 1:
|
||||
return self._retrieve_equity(candidates[0][0])
|
||||
return self._retrieve_equity(candidates[0]['sid'])
|
||||
|
||||
# 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])
|
||||
elif not candidates:
|
||||
sid = sa.select((equities_cols.sid,)).where(
|
||||
(equities_cols.symbol == symbol) &
|
||||
(equities_cols.start_date <= ad_value),
|
||||
).order_by(
|
||||
equities_cols.end_date.desc(),
|
||||
).scalar().execute()
|
||||
if sid:
|
||||
return self._retrieve_equity(sid)
|
||||
|
||||
# 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])
|
||||
elif len(candidates) > 1:
|
||||
sid = sa.select((equities_cols.sid,)).where(
|
||||
(equities_cols.symbol == symbol) &
|
||||
(equities_cols.start_date <= ad_value),
|
||||
).order_by(
|
||||
equities_cols.start_date.desc(),
|
||||
equities_cols.end_date.desc(),
|
||||
).scalar().execute()
|
||||
if sid:
|
||||
return self._retrieve_equity(sid)
|
||||
|
||||
raise SymbolNotFound(symbol=symbol)
|
||||
|
||||
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:
|
||||
sids = sa.select((equities_cols.sid,)).where(
|
||||
equities_cols.symbol == sid,
|
||||
).execute().fetchall()
|
||||
if len(sids) == 1:
|
||||
return self._retrieve_equity(sids[0]['sid'])
|
||||
elif not sids:
|
||||
raise SymbolNotFound(symbol=symbol)
|
||||
else:
|
||||
options = []
|
||||
for row in data:
|
||||
sid = row[0]
|
||||
asset = self._retrieve_equity(sid)
|
||||
options.append(asset)
|
||||
raise MultipleSymbolsFound(symbol=symbol,
|
||||
options=options)
|
||||
raise MultipleSymbolsFound(
|
||||
symbol=symbol,
|
||||
options=list(map(
|
||||
compose(self._retrieve_equity, itemgetter('sid')),
|
||||
sids,
|
||||
))
|
||||
)
|
||||
|
||||
def lookup_symbol(self, symbol, as_of_date, fuzzy=False):
|
||||
def lookup_symbol(self, symbol, as_of_date, fuzzy=None):
|
||||
"""
|
||||
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
|
||||
@@ -291,41 +310,39 @@ class AssetFinder(object):
|
||||
so we can find a match by inserting an underscore.
|
||||
"""
|
||||
symbol = symbol.upper()
|
||||
as_of_date = normalize_date(as_of_date)
|
||||
ad_value = normalize_date(as_of_date).value
|
||||
|
||||
if not fuzzy:
|
||||
if fuzzy is None:
|
||||
try:
|
||||
return self.lookup_symbol_resolve_multiple(symbol, as_of_date)
|
||||
except SymbolNotFound:
|
||||
return None
|
||||
else:
|
||||
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()
|
||||
|
||||
# If one SID exists for symbol, return that symbol
|
||||
if len(candidates) == 1:
|
||||
return self._retrieve_equity(candidates[0][0])
|
||||
equities_cols = self.equities.c
|
||||
candidates = sa.select((equities_cols.sid,)).where(
|
||||
(equities_cols.fuzzy == fuzzy) &
|
||||
(equities_cols.start_date <= ad_value) &
|
||||
(equities_cols.end_date >= ad_value),
|
||||
).execute().fetchall()
|
||||
|
||||
# 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 one SID exists for symbol, return that symbol
|
||||
if len(candidates) == 1:
|
||||
return self._retrieve_equity(candidates[0]['sid'])
|
||||
|
||||
# If multiple SIDs exist for symbol, return latest start_date with
|
||||
# end_date as a tie-breaker
|
||||
elif candidates:
|
||||
sid = sa.select((equities_cols.sid,)).where(
|
||||
(equities_cols.symbol == symbol) &
|
||||
(equities_cols.start_date <= ad_value),
|
||||
).order_by(
|
||||
equities_cols.start_date.desc(),
|
||||
equities_cols.end_date.desc(),
|
||||
).scalar().execute()
|
||||
if sid:
|
||||
return self._retrieve_equity(sid)
|
||||
|
||||
raise SymbolNotFound(symbol=symbol)
|
||||
|
||||
def lookup_future_chain(self, root_symbol, as_of_date, knowledge_date):
|
||||
""" Return the futures chain for a given root symbol.
|
||||
@@ -359,55 +376,54 @@ class AssetFinder(object):
|
||||
Raised when a future chain could not be found for the given
|
||||
root symbol.
|
||||
"""
|
||||
c = self.conn.cursor()
|
||||
|
||||
fc_cols = self.futures_contracts.c
|
||||
|
||||
if as_of_date is pd.NaT:
|
||||
# If the as_of_date is NaT, get all contracts for this
|
||||
# root symbol.
|
||||
t = {'root_symbol': root_symbol}
|
||||
c.execute("""
|
||||
select sid from futures
|
||||
where root_symbol=:root_symbol
|
||||
order by notice_date asc
|
||||
""", t)
|
||||
sids = list(map(
|
||||
itemgetter('sid'),
|
||||
sa.select((fc_cols.sid,)).where(
|
||||
(fc_cols.root_symbol == root_symbol),
|
||||
).order_by(
|
||||
fc_cols.notice_date.asc(),
|
||||
).execute().fetchall()))
|
||||
else:
|
||||
as_of_date = as_of_date.value
|
||||
if knowledge_date is pd.NaT:
|
||||
# If knowledge_date is NaT, default to using as_of_date
|
||||
t = {'root_symbol': root_symbol,
|
||||
'as_of_date': as_of_date.value,
|
||||
'knowledge_date': as_of_date.value}
|
||||
knowledge_date = as_of_date.value
|
||||
else:
|
||||
t = {'root_symbol': root_symbol,
|
||||
'as_of_date': as_of_date.value,
|
||||
'knowledge_date': knowledge_date.value}
|
||||
knowledge_date = knowledge_date.value
|
||||
|
||||
sids = list(map(
|
||||
itemgetter('sid'),
|
||||
sa.select((fc_cols.sid,)).where(
|
||||
(fc_cols.root_symbol == root_symbol) &
|
||||
(fc_cols.notice_date >= as_of_date) &
|
||||
(fc_cols.start_date <= knowledge_date),
|
||||
).order_by(
|
||||
fc_cols.notice_date.asc(),
|
||||
).execute().fetchall()
|
||||
))
|
||||
|
||||
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_contracts
|
||||
WHERE root_symbol=:root_symbol
|
||||
""", t)
|
||||
count = c.fetchone()[0]
|
||||
count = sa.select((sa.func.count(fc_cols.sid),)).where(
|
||||
fc_cols.root_symbol == root_symbol,
|
||||
).scalar().execute()
|
||||
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]
|
||||
|
||||
return map(self._retrieve_futures_contract, sids)
|
||||
|
||||
@property
|
||||
def sids(self):
|
||||
c = self.conn.cursor()
|
||||
query = 'SELECT sid FROM asset_router'
|
||||
c.execute(query)
|
||||
return [r[0] for r in c.fetchall()]
|
||||
return tuple(map(
|
||||
itemgetter('sid'),
|
||||
sa.select(self.asset_router.c.sid,).execute().fetchall(),
|
||||
))
|
||||
|
||||
def _lookup_generic_scalar(self,
|
||||
asset_convertible,
|
||||
@@ -550,33 +566,38 @@ class AssetFinder(object):
|
||||
def _compute_asset_lifetimes(self):
|
||||
"""
|
||||
Compute and cache a recarry of asset lifetimes.
|
||||
|
||||
FUTURE OPTIMIZATION: We're looping over a big array, which means this
|
||||
probably should be in C/Cython.
|
||||
"""
|
||||
with self.conn as transaction:
|
||||
results = transaction.execute(
|
||||
'SELECT sid, start_date, end_date from equities'
|
||||
).fetchall()
|
||||
|
||||
lifetimes = np.recarray(
|
||||
shape=(len(results),),
|
||||
dtype=[('sid', 'i8'), ('start', 'i8'), ('end', 'i8')],
|
||||
)
|
||||
|
||||
# TODO: This is **WAY** slower than it could be because we have to
|
||||
# check for None everywhere. If we represented "no start date" as
|
||||
# 0, and "no end date" as MAX_INT in our metadata, this would be
|
||||
# significantly faster.
|
||||
NO_START = 0
|
||||
NO_END = np.iinfo(int).max
|
||||
for idx, (sid, start, end) in enumerate(results):
|
||||
lifetimes[idx] = (
|
||||
sid,
|
||||
start if start is not None else NO_START,
|
||||
end if end is not None else NO_END,
|
||||
)
|
||||
return lifetimes
|
||||
equities_cols = self.equities.c
|
||||
buf = np.array(
|
||||
tuple(map(
|
||||
float,
|
||||
sa.select((
|
||||
equities_cols.sid,
|
||||
equities_cols.start_date,
|
||||
equities_cols.end_date,
|
||||
)).execute(),
|
||||
)),
|
||||
dype='<f8', # use doubles so we get NaNs
|
||||
)
|
||||
lifetimes = np.recarray(
|
||||
buf=buf,
|
||||
shape=(len(buf),),
|
||||
dtype=[
|
||||
('sid', '<f8'),
|
||||
('start', '<f8'),
|
||||
('end', '<f8')
|
||||
],
|
||||
)
|
||||
start = lifetimes.start
|
||||
end = lifetimes.end
|
||||
start[np.isnan(start)] = 0 # convert missing starts to 0
|
||||
end[np.isnan(end)] = np.iinfo(int).max # convert missing end to INTMAX
|
||||
# Cast the results back down to int.
|
||||
return lifetimes.astype([
|
||||
('sid', '<i8'),
|
||||
('start', '<i8'),
|
||||
('end', '<i8'),
|
||||
])
|
||||
|
||||
def lifetimes(self, dates):
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user