fix update_bundle

rqalpha/__main__.py

@@ -20,8 +20,12 @@
 import csv
 import os
 import click
+import tempfile
+import tarfile
+import datetime
+import requests
 import pandas as pd
-from six import StringIO, iteritems
+from six import StringIO, iteritems, print_
 
 from .cache_control import set_cache_policy, CachePolicy
 from .utils.click_helper import Date
@@ -40,6 +44,41 @@ def entry_point():
     cli(obj={})
 
 
+@cli.command()
+@click.option('-d', '--data-bundle-path', default=os.path.expanduser("~/.rqalpha"), type=click.Path())
+def update_bundle(data_bundle_path):
+    """update data bundle, download if not found"""
+    day = datetime.date.today()
+    tmp = os.path.join(tempfile.gettempdir(), 'rq.bundle')
+
+    while True:
+        url = 'http://7xjci3.com1.z0.glb.clouddn.com/bundles/rqbundle_%04d%02d%02d.tar.bz2' % (
+        day.year, day.month, day.day)
+        print_('try {} ...'.format(url))
+        r = requests.get(url, stream=True)
+        if r.status_code != 200:
+            day = day - datetime.timedelta(days=1)
+            continue
+
+        out = open(tmp, 'wb')
+        total_length = int(r.headers.get('content-length'))
+
+        with click.progressbar(length=total_length, label='downloading ...') as bar:
+            for data in r.iter_content(chunk_size=8192):
+                bar.update(len(data))
+                out.write(data)
+
+        out.close()
+        break
+
+    shutil.rmtree(data_bundle_path, ignore_errors=True)
+    os.mkdir(data_bundle_path)
+    tar = tarfile.open(tmp, 'r:bz2')
+    tar.extractall(data_bundle_path)
+    tar.close()
+    os.remove(tmp)
+
+
 @cli.command()
 @click.option('-d', '--data-bundle-path', 'base__data_bundle_path', type=click.Path(exists=True))
 @click.option('-f', '--strategy-file', 'base__strategy_file', type=click.Path(exists=True))

rqalpha/config.yml

@@ -5,7 +5,7 @@ base:
   # 可以指定回测的唯一ID，用户区分多次回测的结果
   run_id: 9999
   # 数据源所存储的文件路径
-  data_bundle_path: ./bundle/
+  data_bundle_path: ~
   # 启动的策略文件路径
   strategy_file: ./rqalpha/examples/test.py
   # 回测起始日期

rqalpha/utils/config.py

@@ -75,8 +75,11 @@ def parse_config(click_kargs, base_config_path=None):
     if base_config.margin_multiplier <= 0:
         raise patch_user_exc(ValueError(_("invalid margin multiplier value: value range is (0, +∞]")))
 
+    if base_config.data_bundle_path is None:
+        base_config.data_bundle_path = os.path.expanduser("~/.rqalpha")
+
     if not os.path.exists(base_config.data_bundle_path):
-        print("data bundle not found. Run `python3 {baseProject}/tools/create_bundle` to download data bundle.")
+        print("data bundle not found. Run `rqalpha update_bundle` to download data bundle.")
         # print("data bundle not found. Run `%s update_bundle` to download data bundle." % sys.argv[0])
         return