diff --git a/zipline/assets/asset_writer.py b/zipline/assets/asset_writer.py index dfc8566b..d11df1c0 100644 --- a/zipline/assets/asset_writer.py +++ b/zipline/assets/asset_writer.py @@ -346,7 +346,6 @@ class AssetDBWriter(object): -------- zipline.assets.asset_finder """ - with self.engine.begin() as txn: # Create SQL tables if they do not exist. metadata = self.init_db(txn) @@ -358,7 +357,6 @@ class AssetDBWriter(object): exchanges if exchanges is not None else pd.DataFrame(), root_symbols if root_symbols is not None else pd.DataFrame(), ) - # Write the data to SQL. self._write_df_to_table( metadata.tables['futures_exchanges'], @@ -468,7 +466,6 @@ class AssetDBWriter(object): check_version_info(version_info, ASSET_DB_VERSION) else: write_version_info(version_info, ASSET_DB_VERSION) - return metadata def _normalize_equities(self, equities): diff --git a/zipline/data/bundles/core.py b/zipline/data/bundles/core.py index a0765f54..89064577 100644 --- a/zipline/data/bundles/core.py +++ b/zipline/data/bundles/core.py @@ -22,7 +22,6 @@ from ..minute_bars import ( from zipline.assets import AssetDBWriter, AssetFinder, ASSET_DB_VERSION from zipline.utils.cache import ( dataframe_cache, - working_file, working_dir, ) from zipline.utils.compat import mappingproxy @@ -38,39 +37,59 @@ open_and_closes = nyse_cal.schedule def asset_db_path(bundle_name, timestr, environ=None): return pth.data_path( - [bundle_name, timestr, 'assets-%d.sqlite' % ASSET_DB_VERSION], + asset_db_relative(bundle_name, timestr, environ), environ=environ, ) def minute_equity_path(bundle_name, timestr, environ=None): return pth.data_path( - [bundle_name, timestr, 'minute_equities.bcolz'], + minute_equity_relative(bundle_name, timestr, environ), environ=environ, ) def daily_equity_path(bundle_name, timestr, environ=None): return pth.data_path( - [bundle_name, timestr, 'daily_equities.bcolz'], + daily_equity_relative(bundle_name, timestr, environ), environ=environ, ) def adjustment_db_path(bundle_name, timestr, environ=None): return pth.data_path( - [bundle_name, timestr, 'adjustments.sqlite'], + adjustment_db_relative(bundle_name, timestr, environ), environ=environ, ) def cache_path(bundle_name, environ=None): return pth.data_path( - [bundle_name, '.cache'], + cache_relative(bundle_name, environ), environ=environ, ) +def adjustment_db_relative(bundle_name, timestr, environ=None): + return bundle_name, timestr, 'adjustments.sqlite' + + +def cache_relative(bundle_name, timestr, environ=None): + return bundle_name, '.cache' + + +def daily_equity_relative(bundle_name, timestr, environ=None): + return bundle_name, timestr, 'daily_equities.bcolz' + + +def minute_equity_relative(bundle_name, timestr, environ=None): + return bundle_name, timestr, 'minute_equities.bcolz' + + +def asset_db_relative(bundle_name, timestr, environ=None): + return bundle_name, timestr, 'assets-%d.sqlite' % ASSET_DB_VERSION + + def to_bundle_ingest_dirname(ts): """Convert a pandas Timestamp into the name of the directory for the ingestion. @@ -321,16 +340,19 @@ def _make_bundle_core(): cachepath = cache_path(name, environ=environ) pth.ensure_directory(pth.data_path([name, timestr], environ=environ)) pth.ensure_directory(cachepath) - with dataframe_cache(cachepath, clean_on_failure=False) as cache, \ ExitStack() as stack: # we use `cleanup_on_failure=False` so that we don't purge the # cache directory if the load fails in the middle - if bundle.create_writers: - daily_bars_path = stack.enter_context(working_dir( - daily_equity_path(name, timestr, environ=environ), - )).path + wd = stack.enter_context(working_dir( + pth.data_path([], environ=environ)) + ) + daily_bars_path = wd.ensure_dir( + *daily_equity_relative( + name, timestr, environ=environ, + ) + ) daily_bar_writer = BcolzDailyBarWriter( daily_bars_path, nyse_cal, @@ -341,37 +363,37 @@ def _make_bundle_core(): # when we create the SQLiteAdjustmentWriter below. The # SQLiteAdjustmentWriter needs to open the daily ctables so # that it can compute the adjustment ratios for the dividends. + daily_bar_writer.write(()) minute_bar_writer = BcolzMinuteBarWriter( bundle.calendar[0], - stack.enter_context(working_dir( - minute_equity_path(name, timestr, environ=environ), - )).path, + wd.ensure_dir(*minute_equity_relative( + name, timestr, environ=environ) + ), bundle.opens, bundle.closes, minutes_per_day=bundle.minutes_per_day, ) asset_db_writer = AssetDBWriter( - stack.enter_context(working_file( - asset_db_path(name, timestr, environ=environ), - )).path, + wd.getpath(*asset_db_relative( + name, timestr, environ=environ, + )) ) - wf = stack.enter_context(working_file( - adjustment_db_path(name, timestr, environ=environ), - )) - wf.close() # we need to close the file to delete it on windows - adjustment_db_writer = SQLiteAdjustmentWriter( - wf.path, - BcolzDailyBarReader(daily_bars_path), - bundle.calendar, - overwrite=True, + + adjustment_db_writer = stack.enter_context( + SQLiteAdjustmentWriter( + wd.getpath(*adjustment_db_relative( + name, timestr, environ=environ)), + BcolzDailyBarReader(daily_bars_path), + bundle.calendar, + overwrite=True, + ) ) else: daily_bar_writer = None minute_bar_writer = None asset_db_writer = None adjustment_db_writer = None - bundle.ingest( environ, asset_db_writer, @@ -541,5 +563,4 @@ def _make_bundle_core(): return BundleCore(bundles, register, unregister, ingest, load, clean) - bundles, register, unregister, ingest, load, clean = _make_bundle_core() diff --git a/zipline/data/us_equity_pricing.py b/zipline/data/us_equity_pricing.py index cd94512b..a6d92504 100644 --- a/zipline/data/us_equity_pricing.py +++ b/zipline/data/us_equity_pricing.py @@ -210,8 +210,15 @@ class BcolzDailyBarWriter(object): def __init__(self, filename, calendar, start_session, end_session): self._filename = filename - assert calendar.is_session(start_session), "Start session is invalid!" - assert calendar.is_session(end_session), "End session is invalid!" + if start_session != end_session: + if not calendar.is_session(start_session): + raise ValueError( + "Start session %s is invalid!" % start_session + ) + if not calendar.is_session(end_session): + raise ValueError( + "End session %s is invalid!" % end_session + ) self._start_session = start_session self._end_session = end_session @@ -1082,6 +1089,12 @@ class SQLiteAdjustmentWriter(object): dividend_ratios = self.calc_dividend_ratios(dividends) self.write_frame('dividends', dividend_ratios) + def __enter__(self): + return self + + def __exit__(self, *exc_info): + self.close() + def write(self, splits=None, mergers=None, diff --git a/zipline/utils/cache.py b/zipline/utils/cache.py index d2a9b14e..21aee7d4 100644 --- a/zipline/utils/cache.py +++ b/zipline/utils/cache.py @@ -5,7 +5,8 @@ from collections import namedtuple, MutableMapping import errno import os import pickle -from shutil import rmtree, copyfile, copytree +from distutils import dir_util +from shutil import rmtree, move from tempfile import mkdtemp, NamedTemporaryFile import pandas as pd @@ -272,11 +273,11 @@ class working_file(object): Notes ----- The file is moved on __exit__ if there are no exceptions. - ``working_file`` uses :func:`shutil.copyfile` to move the actual files, - meaning it has as strong of guarantees as :func:`shutil.copyfile`. + ``working_file`` uses :func:`shutil.move` to move the actual files, + meaning it has as strong of guarantees as :func:`shutil.move`. """ def __init__(self, final_path, *args, **kwargs): - self._tmpfile = NamedTemporaryFile(*args, **kwargs) + self._tmpfile = NamedTemporaryFile(delete=False, *args, **kwargs) self._final_path = final_path @property @@ -289,7 +290,8 @@ class working_file(object): def _commit(self): """Sync the temporary file to the final path. """ - copyfile(self.name, self._final_path) + self._tmpfile.close() + move(self._name, self._final_path) def __getattr__(self, attr): return getattr(self._tmpfile, attr) @@ -299,9 +301,9 @@ class working_file(object): return self def __exit__(self, *exc_info): + self._tmpfile.__exit__(*exc_info) if exc_info[0] is None: self._commit() - self._tmpfile.__exit__(*exc_info) class working_dir(object): @@ -318,15 +320,15 @@ class working_dir(object): Notes ----- The file is moved on __exit__ if there are no exceptions. - ``working_dir`` uses :func:`shutil.copytree` to move the actual files, - meaning it has as strong of guarantees as :func:`shutil.copytree`. + ``working_dir`` uses :func:`dir_util.copy_tree` to move the actual files, + meaning it has as strong of guarantees as :func:`dir_util.copy_tree`. """ def __init__(self, final_path, *args, **kwargs): self.path = mkdtemp() self._final_path = final_path - def mkdir(self, *path_parts): - """Create a subdirectory of the working directory. + def ensure_dir(self, *path_parts): + """Ensures a subdirectory of the working directory. Parameters ---------- @@ -334,7 +336,7 @@ class working_dir(object): The parts of the path after the working directory. """ path = self.getpath(*path_parts) - os.mkdir(path) + ensure_directory(path) return path def getpath(self, *path_parts): @@ -350,7 +352,7 @@ class working_dir(object): def _commit(self): """Sync the temporary directory to the final path. """ - copytree(self.path, self._final_path) + dir_util.copy_tree(self.path, self._final_path) def __enter__(self): return self