From 0b3a35891ec9ca6f7e63f0d83e43dba0ae9e5e9f Mon Sep 17 00:00:00 2001 From: Joe Jevnik Date: Thu, 5 May 2016 15:45:23 -0400 Subject: [PATCH] ENH: fix the quality of life issues in the CLI Fixes the issues presented in #1181 by @ssanderson around the new command line interface. --- setup.py | 5 +++++ zipline/__main__.py | 38 +++++++++++++++++++++++++++------- zipline/assets/assets.py | 2 +- zipline/data/bundles/quandl.py | 2 ++ zipline/utils/cache.py | 14 ++++++------- zipline/utils/run_algo.py | 2 +- 6 files changed, 45 insertions(+), 18 deletions(-) diff --git a/setup.py b/setup.py index 78063820..16f903d2 100644 --- a/setup.py +++ b/setup.py @@ -260,6 +260,11 @@ setup( version=versioneer.get_version(), cmdclass=LazyBuildExtCommandClass(versioneer.get_cmdclass()), description='A backtester for financial algorithms.', + entry_points={ + 'console_scripts': [ + 'zipline = zipline.__main__:main', + ], + }, author='Quantopian Inc.', author_email='opensource@quantopian.com', packages=find_packages('.', include=['zipline', 'zipline.*']), diff --git a/zipline/__main__.py b/zipline/__main__.py index 7afd2c1e..db496e86 100644 --- a/zipline/__main__.py +++ b/zipline/__main__.py @@ -5,8 +5,9 @@ import click import logbook import pandas as pd -from zipline.data import bundles +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: @@ -36,7 +37,7 @@ except NameError: default=True, help="Don't load the default zipline extension.py file in $ZIPLINE_HOME.", ) -def cli(extension, strict_extensions, default_extension): +def main(extension, strict_extensions, default_extension): """Top level zipline entry point. """ # install a logbook handler before performing any other operations @@ -97,7 +98,7 @@ def ipython_only(option): return d -@cli.command() +@main.command() @click.option( '-f', '--algofile', @@ -273,7 +274,7 @@ def zipline_magic(line, cell=None): raise ValueError('main returned non-zero status code: %d' % e.code) -@cli.command() +@main.command() @click.option( '-b', '--bundle', @@ -291,7 +292,7 @@ def zipline_magic(line, cell=None): def ingest(bundle, show_progress): """Ingest the data for the given bundle. """ - bundles.ingest( + bundles_module.ingest( bundle, os.environ, pd.Timestamp.utcnow(), @@ -299,7 +300,7 @@ def ingest(bundle, show_progress): ) -@cli.command() +@main.command() @click.option( '-b', '--bundle', @@ -333,7 +334,7 @@ def ingest(bundle, show_progress): def clean(bundle, before, after, keep_last): """Clean up data downloaded with the ingest command. """ - bundles.clean( + bundles_module.clean( bundle, before, after, @@ -341,5 +342,26 @@ def clean(bundle, before, after, keep_last): ) +@main.command() +def bundles(): + """List all of the available data bundles. + """ + for bundle in sorted(bundles_module.bundles.keys()): + ingestions = sorted( + (str(pd.Timestamp(int(ing))) + for ing in os.listdir(pth.data_path([bundle])) + if not pth.hidden(ing)), + reverse=True, + ) + print( + '\n'.join( + '%s %s' % (bundle, line) + for line in ( + ingestions if ingestions else ('',) + ) + ), + ) + + if __name__ == '__main__': - cli() + main() diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index b671a754..f9ffc934 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -101,7 +101,7 @@ class AssetFinder(object): PERSISTENT_TOKEN = "" def __init__(self, engine): - if isinstance(engine, str): + if isinstance(engine, string_types): engine = sa.create_engine('sqlite:///' + engine) self.engine = engine diff --git a/zipline/data/bundles/quandl.py b/zipline/data/bundles/quandl.py index a41462d5..708e9f35 100644 --- a/zipline/data/bundles/quandl.py +++ b/zipline/data/bundles/quandl.py @@ -319,6 +319,8 @@ def quantopian_quandl_bundle(environ, cache, show_progress, output_dir): + if show_progress: + print('Downloading quandl data, can take about 1 minute') # use closing for py2 compat with closing(urlopen(QUANTOPIAN_QUANDL_URL)) as f, \ tarfile.open('r', fileobj=BytesIO(f.read())) as tar: diff --git a/zipline/utils/cache.py b/zipline/utils/cache.py index 15dff318..f272524b 100644 --- a/zipline/utils/cache.py +++ b/zipline/utils/cache.py @@ -194,7 +194,7 @@ class dataframe_cache(MutableMapping): self._protocol = int(s[1]) if len(s) == 2 else None self.serialize = self._serialize_pickle - self.deserialize = self._deserialize_pickle + self.deserialize = pickle.load ensure_directory(self.path) @@ -202,10 +202,6 @@ class dataframe_cache(MutableMapping): with open(path, 'wb') as f: pickle.dump(df, f, protocol=self._protocol) - def _deserialize_pickle(self, path): - with open(path, 'rb') as f: - return pickle.load(f) - def _keypath(self, key): return os.path.join(self.path, key) @@ -226,9 +222,11 @@ class dataframe_cache(MutableMapping): with self.lock: try: - return self.deserialize(self._keypath(key)) - except UnboundLocalError: - # This is how pandas fails if the file doesn't exist! #pandas + with open(self._keypath(key), 'rb') as f: + return self.deserialize(f) + except IOError as e: + if e.errno != errno.ENOENT: + raise raise KeyError(key) def __setitem__(self, key, value): diff --git a/zipline/utils/run_algo.py b/zipline/utils/run_algo.py index 4a8854b7..288307ba 100644 --- a/zipline/utils/run_algo.py +++ b/zipline/utils/run_algo.py @@ -145,7 +145,7 @@ def _run(handle_data, 'before_trading_start': before_trading_start, 'analyze': analyze, } if algotext is None else { - 'algo_filename': algofile, + 'algo_filename': algofile.name, 'script': algotext, } ).run(