mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-13 05:14:54 +08:00
ENH: Makes working_dir work on Windows. Updates bundle core
This commit is contained in:
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
+14
-12
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user