From ee4aa7327b6243243f9c85e4c60b1516ccc439f5 Mon Sep 17 00:00:00 2001 From: llllllllll Date: Thu, 6 Aug 2015 12:06:33 -0400 Subject: [PATCH] MAINT: more bugfixes --- zipline/assets/asset_writer.py | 26 +++++++++++++++++++------- zipline/assets/assets.py | 12 ++++++------ zipline/finance/trading.py | 16 ++++++++-------- 3 files changed, 33 insertions(+), 21 deletions(-) diff --git a/zipline/assets/asset_writer.py b/zipline/assets/asset_writer.py index 06e19a6e..151ea576 100644 --- a/zipline/assets/asset_writer.py +++ b/zipline/assets/asset_writer.py @@ -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), diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index 827d3447..ae05b70c 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -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) diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index cda614ed..07170f03 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -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')