wassname/catalyst
1
0
Fork 0
mirror of https://github.com/wassname/catalyst.git synced 2026-07-11 08:03:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

DEV: zipline ingest can downgrade the assets db

Browse Source 
This lets us publish an "old" db version for the most recent release
of zipline, using the latest code base.
This commit is contained in:
Richard Frank
2016-08-23 12:16:50 -04:00
parent 418e8901f4
commit 419bd1e3b5
5 changed files with 143 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/data/bundles/test_core.py
+92 -15
View File
@@ -2,12 +2,17 @@ import os
from nose_parameterized import parameterized
import pandas as pd
import sqlalchemy as sa
from toolz import valmap
import toolz.curried.operator as op
from zipline.assets import ASSET_DB_VERSION
from zipline.assets.asset_writer import check_version_info
from zipline.assets.synthetic import make_simple_equity_info
from zipline.data.bundles import UnknownBundle, from_bundle_ingest_dirname
from zipline.data.bundles.core import _make_bundle_core, BadClean
from zipline.data.bundles import UnknownBundle, from_bundle_ingest_dirname, \
ingestions_for_bundle
from zipline.data.bundles.core import _make_bundle_core, BadClean, \
to_bundle_ingest_dirname, asset_db_path
from zipline.lib.adjustment import Float64Multiply
from zipline.pipeline.loaders.synthetic import (
make_bar_data,
@@ -17,7 +22,8 @@ from zipline.testing import (
subtest,
str_to_seconds,
)
from zipline.testing.fixtures import WithInstanceTmpDir, ZiplineTestCase
from zipline.testing.fixtures import WithInstanceTmpDir, ZiplineTestCase, \
WithDefaultDateBounds
from zipline.testing.predicates import (
assert_equal,
assert_false,
@@ -37,7 +43,13 @@ import zipline.utils.paths as pth
_1_ns = pd.Timedelta(1, unit='ns')
class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
class BundleCoreTestCase(WithInstanceTmpDir,
WithDefaultDateBounds,
ZiplineTestCase):
START_DATE = pd.Timestamp('2014-01-06', tz='utc')
END_DATE = pd.Timestamp('2014-01-10', tz='utc')
def init_instance_fixtures(self):
super(BundleCoreTestCase, self).init_instance_fixtures()
(self.bundles,
@@ -111,18 +123,17 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
assert_true(called[0])
def test_ingest(self):
start = pd.Timestamp('2014-01-06', tz='utc')
end = pd.Timestamp('2014-01-10', tz='utc')
calendar = get_calendar('NYSE')
sessions = calendar.sessions_in_range(start, end)
minutes = calendar.minutes_for_sessions_in_range(start, end)
sessions = calendar.sessions_in_range(self.START_DATE, self.END_DATE)
minutes = calendar.minutes_for_sessions_in_range(
self.START_DATE, self.END_DATE,
)
sids = tuple(range(3))
equities = make_simple_equity_info(
sids,
start,
end,
self.START_DATE,
self.END_DATE,
)
daily_bar_data = make_bar_data(equities, sessions)
@@ -145,8 +156,8 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
@self.register(
'bundle',
calendar=calendar,
start_session=start,
end_session=end,
start_session=self.START_DATE,
end_session=self.END_DATE,
)
def bundle_ingest(environ,
asset_db_writer,
@@ -193,8 +204,8 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
actual = bundle.equity_daily_bar_reader.load_raw_arrays(
columns,
start,
end,
self.START_DATE,
self.END_DATE,
sids,
)
for actual_column, colname in zip(actual, columns):
@@ -253,6 +264,72 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase):
msg='volume',
)
def test_ingest_assets_versions(self):
versions = (1, 2)
called = [False]
@self.register('bundle', create_writers=False)
def bundle_ingest_no_create_writers(*args, **kwargs):
called[0] = True
now = pd.Timestamp.utcnow()
with self.assertRaisesRegexp(ValueError,
"ingest .* creates writers .* downgrade"):
self.ingest('bundle', self.environ, assets_versions=versions,
timestamp=now - pd.Timedelta(seconds=1))
assert_false(called[0])
assert_equal(len(ingestions_for_bundle('bundle', self.environ)), 1)
@self.register('bundle', create_writers=True)
def bundle_ingest_create_writers(
environ,
asset_db_writer,
minute_bar_writer,
daily_bar_writer,
adjustment_writer,
calendar,
start_session,
end_session,
cache,
show_progress,
output_dir):
self.assertIsNotNone(asset_db_writer)
self.assertIsNotNone(minute_bar_writer)
self.assertIsNotNone(daily_bar_writer)
self.assertIsNotNone(adjustment_writer)
equities = make_simple_equity_info(
tuple(range(3)),
self.START_DATE,
self.END_DATE,
)
asset_db_writer.write(equities=equities)
called[0] = True
# Explicitly use different timestamp; otherwise, test could run so fast
# that first ingestion is re-used.
self.ingest('bundle', self.environ, assets_versions=versions,
timestamp=now)
assert_true(called[0])
ingestions = ingestions_for_bundle('bundle', self.environ)
assert_equal(len(ingestions), 2)
for version in sorted(set(versions) | {ASSET_DB_VERSION}):
eng = sa.create_engine(
'sqlite:///' +
asset_db_path(
'bundle',
to_bundle_ingest_dirname(ingestions[0]), # most recent
self.environ,
version,
)
)
metadata = sa.MetaData()
metadata.reflect(eng)
version_table = metadata.tables['version_info']
check_version_info(eng, version_table, version)
@parameterized.expand([('clean',), ('load',)])
def test_bundle_doesnt_exist(self, fnname):
with assert_raises(UnknownBundle) as e:
zipline/__main__.py
+12 -9
View File
@@ -5,10 +5,10 @@ from functools import wraps
import click
import logbook
import pandas as pd
from six import text_type
from zipline.data import bundles as bundles_module
from zipline.utils.cli import Date, Timestamp
import zipline.utils.paths as pth
from zipline.utils.run_algo import _run, load_extensions
try:
@@ -284,19 +284,25 @@ def zipline_magic(line, cell=None):
show_default=True,
help='The data bundle to ingest.',
)
@click.option(
'--assets-version',
type=int,
multiple=True,
help='Version of the assets db to which to downgrade.',
)
@click.option(
'--show-progress/--no-show-progress',
is_flag=True,
default=True,
help='Print progress information to the terminal.'
)
def ingest(bundle, show_progress):
def ingest(bundle, assets_version, show_progress):
"""Ingest the data for the given bundle.
"""
bundles_module.ingest(
bundle,
os.environ,
pd.Timestamp.utcnow(),
assets_version,
show_progress,
)
@@ -352,11 +358,8 @@ def bundles():
# hide the test data
continue
try:
ingestions = sorted(
(str(bundles_module.from_bundle_ingest_dirname(ing))
for ing in os.listdir(pth.data_path([bundle]))
if not pth.hidden(ing)),
reverse=True,
ingestions = list(
map(text_type, bundles_module.ingestions_for_bundle(bundle))
)
except OSError as e:
if e.errno != errno.ENOENT:
@@ -367,7 +370,7 @@ def bundles():
# because there were no entries, print a single message indicating that
# no ingestions have yet been made.
for timestamp in ingestions or ["<no ingestions>"]:
print("%s %s" % (bundle, timestamp))
click.echo("%s %s" % (bundle, timestamp))
if __name__ == '__main__':
zipline/data/bundles/__init__.py
+2
View File
@@ -6,6 +6,7 @@ from .core import (
clean,
from_bundle_ingest_dirname,
ingest,
ingestions_for_bundle,
load,
register,
to_bundle_ingest_dirname,
@@ -20,6 +21,7 @@ __all__ = [
'clean',
'from_bundle_ingest_dirname',
'ingest',
'ingestions_for_bundle',
'load',
'register',
'to_bundle_ingest_dirname',
zipline/data/bundles/core.py
+36 -9
View File
@@ -20,9 +20,11 @@ from ..minute_bars import (
BcolzMinuteBarWriter,
)
from zipline.assets import AssetDBWriter, AssetFinder, ASSET_DB_VERSION
from zipline.assets.asset_db_migrations import downgrade
from zipline.utils.cache import (
dataframe_cache,
working_dir,
working_file,
)
from zipline.utils.compat import mappingproxy
from zipline.utils.input_validation import ensure_timestamp, optionally
@@ -31,9 +33,9 @@ from zipline.utils.preprocess import preprocess
from zipline.utils.calendars import get_calendar, register_calendar
def asset_db_path(bundle_name, timestr, environ=None):
def asset_db_path(bundle_name, timestr, environ=None, db_version=None):
return pth.data_path(
asset_db_relative(bundle_name, timestr, environ),
asset_db_relative(bundle_name, timestr, environ, db_version),
environ=environ,
)
@@ -82,8 +84,10 @@ 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 asset_db_relative(bundle_name, timestr, environ=None, db_version=None):
db_version = ASSET_DB_VERSION if db_version is None else db_version
return bundle_name, timestr, 'assets-%d.sqlite' % db_version
def to_bundle_ingest_dirname(ts):
@@ -119,6 +123,15 @@ def from_bundle_ingest_dirname(cs):
return pd.Timestamp(cs.replace(';', ':'))
def ingestions_for_bundle(bundle, environ=None):
return sorted(
(from_bundle_ingest_dirname(ing)
for ing in os.listdir(pth.data_path([bundle], environ))
if not pth.hidden(ing)),
reverse=True,
)
_BundlePayload = namedtuple(
'_BundlePayload',
'calendar start_session end_session minutes_per_day ingest create_writers',
@@ -328,6 +341,7 @@ def _make_bundle_core():
def ingest(name,
environ=os.environ,
timestamp=None,
assets_versions=(),
show_progress=False):
"""Ingest data for a given bundle.
@@ -340,6 +354,8 @@ def _make_bundle_core():
timestamp : datetime, optional
The timestamp to use for the load.
By default this is the current time.
assets_versions : Iterable[int], optional
Versions of the assets db to which to downgrade.
show_progress : bool, optional
Tell the ingest function to display the progress where possible.
"""
@@ -389,11 +405,10 @@ def _make_bundle_core():
bundle.end_session,
minutes_per_day=bundle.minutes_per_day,
)
asset_db_writer = AssetDBWriter(
wd.getpath(*asset_db_relative(
name, timestr, environ=environ,
))
)
assets_db_path = wd.getpath(*asset_db_relative(
name, timestr, environ=environ,
))
asset_db_writer = AssetDBWriter(assets_db_path)
adjustment_db_writer = stack.enter_context(
SQLiteAdjustmentWriter(
@@ -409,6 +424,10 @@ def _make_bundle_core():
minute_bar_writer = None
asset_db_writer = None
adjustment_db_writer = None
if assets_versions:
raise ValueError('Need to ingest a bundle that creates '
'writers in order to downgrade the assets'
' db.')
bundle.ingest(
environ,
asset_db_writer,
@@ -423,6 +442,14 @@ def _make_bundle_core():
pth.data_path([name, timestr], environ=environ),
)
for version in sorted(set(assets_versions), reverse=True):
version_path = wd.getpath(*asset_db_relative(
name, timestr, environ=environ, db_version=version,
))
with working_file(version_path) as wf:
shutil.copy2(assets_db_path, wf.path)
downgrade(wf.path, version)
def most_recent_data(bundle_name, timestamp, environ=None):
"""Get the path to the most recent data after ``date``for the
given bundle.
zipline/utils/cache.py
+1 -5
View File
@@ -290,11 +290,7 @@ class working_file(object):
def _commit(self):
"""Sync the temporary file to the final path.
"""
self._tmpfile.close()
move(self._name, self._final_path)
def __getattr__(self, attr):
return getattr(self._tmpfile, attr)
move(self.path, self._final_path)
def __enter__(self):
self._tmpfile.__enter__()