MAINT: more bugfixes

This commit is contained in:
llllllllll
2015-08-06 12:06:33 -04:00
committed by jfkirk
parent 9cd3707440
commit ee4aa7327b
3 changed files with 33 additions and 21 deletions
+19 -7
View File
@@ -121,23 +121,32 @@ class AssetDBWriter(with_metaclass(ABCMeta)):
self._write_exchanges(exchanges, txn)
self._write_root_symbols(root_symbols, txn)
self._write_futures(futures, txn)
self._write_equities(equities, txn)
self._write_equities(equities, fuzzy_char, txn)
def _write_exchanges(self, exchanges, bind=None):
self.futures_exchanges.insert().values(
exchanges.to_records(),
exchanges.reset_index().rename_axis(
{'index': 'exchange_id'},
1,
).to_dict('records'),
).execute(bind=bind)
def _write_root_symbols(self, root_symbols, bind=None):
self.futures_root_symbols.insert().values(
root_symbols.to_records(),
root_symbols.reset_index().rename_axis(
{'index': 'root_symbol_id'},
1,
).to_dict('records'),
).execute(bind=bind)
def _write_futures(self, futures, bind=None):
recs = futures.to_records()
recs = futures.reset_index().rename_axis(
{'index': 'sid'},
1,
).to_dict('records')
self.futures_contracts.insert().values(recs).execute(bind=bind)
self.asset_router.insert().values([
(rec['sid'], 'future') for rec in recs
(rec['index'], 'future') for rec in recs
]).execute(bind=bind)
def _write_equities(self, equities, fuzzy_char, bind=None):
@@ -145,7 +154,10 @@ class AssetDBWriter(with_metaclass(ABCMeta)):
if fuzzy_char:
equities['fuzzy'] = equities['symbol'].str.replace(fuzzy_char, '')
recs = equities.to_records()
recs = equities.reset_index().rename_axis(
{'index': 'sid'},
1,
).to_dict('records')
self.equities.insert().values(recs).execute(bind=bind)
self.asset_router.insert().values([
(rec['sid'], 'equity') for rec in recs
@@ -233,7 +245,7 @@ class AssetDBWriter(with_metaclass(ABCMeta)):
*((sa.ForeignKey(self.futures_exchanges.c.exchange_id),)
if constraints else ())
),
sa.column('exchange', sa.Text),
sa.Column('exchange', sa.Text),
sa.Column('notice_date', sa.Integer),
sa.Column('expiration_date', sa.Integer),
sa.Column('contract_multiplier', sa.Float),
+6 -6
View File
@@ -52,7 +52,7 @@ class AssetFinder(object):
self.allow_sid_assignment = allow_sid_assignment
self.engine = engine
metadata = sa.Metadata(bind=engine)
metadata = sa.MetaData(bind=engine)
self.equities = equities = sa.Table(
'equities',
metadata,
@@ -131,7 +131,7 @@ class AssetFinder(object):
asset_type = sa.select((self.asset_router.c.asset_type,)).where(
self.asset_router.c.sid == int(sid),
).scalar().execute()
).scalar()
if asset_type is not None:
self._asset_type_cache[sid] = asset_type
@@ -264,7 +264,7 @@ class AssetFinder(object):
(equities_cols.start_date <= ad_value),
).order_by(
equities_cols.end_date.desc(),
).scalar().execute()
).scalar()
if sid:
return self._retrieve_equity(sid)
@@ -277,7 +277,7 @@ class AssetFinder(object):
).order_by(
equities_cols.start_date.desc(),
equities_cols.end_date.desc(),
).scalar().execute()
).scalar()
if sid:
return self._retrieve_equity(sid)
@@ -338,7 +338,7 @@ class AssetFinder(object):
).order_by(
equities_cols.start_date.desc(),
equities_cols.end_date.desc(),
).scalar().execute()
).scalar()
if sid:
return self._retrieve_equity(sid)
@@ -412,7 +412,7 @@ class AssetFinder(object):
# Check if root symbol exists.
count = sa.select((sa.func.count(fc_cols.sid),)).where(
fc_cols.root_symbol == root_symbol,
).scalar().execute()
).scalar()
if count == 0:
raise RootSymbolNotFound(root_symbol=root_symbol)
+8 -8
View File
@@ -21,6 +21,7 @@ import sqlite3
import pandas as pd
import numpy as np
from sqlalchemy import create_engine
from zipline.data.loader import load_market_data
from zipline.utils import tradingcalendar
@@ -137,10 +138,9 @@ class TradingEnvironment(object):
self.exchange_tz = exchange_tz
self.conn = sqlite3.connect(':memory:')
asset_writer = AssetDBWriterFromDictionary()
asset_writer.write_all(self.conn)
self.asset_finder = AssetFinder(self.conn)
self.engine = engine = create_engine('sqlite:///:memory:')
AssetDBWriterFromDictionary().write_all(engine)
self.asset_finder = AssetFinder(engine)
def __enter__(self, *args, **kwargs):
global environment
@@ -182,7 +182,7 @@ class TradingEnvironment(object):
:return:
"""
if clear_metadata:
self.conn = sqlite3.connect(':memory:')
self.engine = create_engine('sqlite:///:memory:')
if asset_finder is not None:
if not isinstance(asset_finder, AssetFinder):
@@ -190,18 +190,18 @@ class TradingEnvironment(object):
self.asset_finder = asset_finder
if asset_metadata is not None:
self.conn = sqlite3.connect(':memory:')
self.engine = create_engine('sqlite:///:memory:')
if isinstance(asset_metadata, dict):
asset_writer = AssetDBWriterFromDictionary(
equities=asset_metadata)
elif isinstance(asset_metadata, pd.DataFrame):
asset_writer = AssetDBWriterFromDataFrame(
equities=asset_metadata)
asset_writer.write_all(self.conn)
asset_writer.write_all(self.engine)
if identifiers is not None:
asset_writer = AssetDBWriterFromList(equities=identifiers)
asset_writer.write_all(self.conn)
asset_writer.write_all(self.engine)
def normalize_date(self, test_date):
test_date = pd.Timestamp(test_date, tz='UTC')