mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-10 15:17:00 +08:00
MAINT: more bugfixes
This commit is contained in:
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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')
|
||||
|
||||
Reference in New Issue
Block a user