update downloader

paddle_hub/module/module.py

@@ -17,7 +17,7 @@ from __future__ import division
 from __future__ import print_function
 from paddle_hub.tools import utils
 from paddle_hub.tools.logger import logger
-from paddle_hub.tools import downloader
+from paddle_hub.tools.downloader import default_downloader
 from paddle_hub.tools import paddle_helper
 from paddle_hub.module import module_desc_pb2
 from paddle_hub.module.signature import Signature, create_signature
@@ -116,7 +116,8 @@ class Module:
 
     def _init_with_url(self, url):
         utils.check_url_valid(url)
-        module_dir = downloader.download_and_uncompress(module_url)
+        module_dir = default_downloader.download_file_and_uncompress(
+            module_url, save_path=".")
         self._init_with_module_file(module_dir)
 
     def _dump_processor(self):

paddle_hub/tools/downloader.py

@@ -23,11 +23,13 @@ import hashlib
 import requests
 import tempfile
 import tarfile
+from paddle_hub.tools import utils
+from paddle_hub.tools.logger import logger
 
-__all__ = ['MODULE_HOME', 'download', 'md5file', 'download_and_uncompress']
+__all__ = ['MODULE_HOME', 'downloader', 'md5file', 'Downloader']
 
 # TODO(ZeyuChen) add environment varialble to set MODULE_HOME
-MODULE_HOME = os.path.expanduser('~/.cache/paddle/module')
+MODULE_HOME = os.path.expanduser('~/.hub/module')
 
 
 # When running unit tests, there could be multiple processes that
@@ -53,29 +55,29 @@ def md5file(fname):
     return hash_md5.hexdigest()
 
 
-def download_and_uncompress(url, save_name=None):
-    module_name = url.split("/")[-2]
-    dirname = os.path.join(MODULE_HOME, module_name)
-    if not os.path.exists(dirname):
-        os.makedirs(dirname)
-
-    #TODO(ZeyuChen) add download md5 file to verify file completeness
-    file_name = os.path.join(
-        dirname,
-        url.split('/')[-1] if save_name is None else save_name)
+class Downloader:
+    def __init__(self, module_home=None):
+        self.module_home = module_home if module_home else MODULE_HOME
 
-    retry = 0
-    retry_limit = 3
+    def download_file(self, url, save_path=None, save_name=None, retry_limit=3):
+        module_name = url.split("/")[-2]
+        save_path = self.module_home if save_path is None else save_path
+        if not os.path.exists(save_path):
+            utils.mkdir(save_path)
+        save_name = url.split('/')[-1] if save_name is None else save_name
+        file_name = os.path.join(save_path, save_name)
+        retry_times = 0
         while not (os.path.exists(file_name)):
             if os.path.exists(file_name):
-            print("file md5", md5file(file_name))
-        if retry < retry_limit:
-            retry += 1
+                logger.info("file md5", md5file(file_name))
+            if retry_times < retry_limit:
+                retry_times += 1
             else:
                 raise RuntimeError(
                     "Cannot download {0} within retry limit {1}".format(
                         url, retry_limit))
-        print("Cache file %s not found, downloading %s" % (file_name, url))
+            logger.info(
+                "Cache file %s not found, downloading %s" % (file_name, url))
             r = requests.get(url, stream=True)
             total_length = r.headers.get('content-length')
 
@@ -95,22 +97,35 @@ def download_and_uncompress(url, save_name=None):
                             "\r[%s%s]" % ('=' * done, ' ' * (50 - done)))
                         sys.stdout.flush()
 
-    print("file download completed!", file_name)
-    #TODO(ZeyuChen) add md5 check error and file incompleted error, then raise
-    # them and catch them
-    with tarfile.open(file_name, "r:gz") as tar:
+        logger.info("file %s download completed!" % (file_name))
+        return file_name
+
+    def uncompress(self, file, dirname=None, delete_file=False):
+        dirname = os.path.dirname(file) if dirname is None else dirname
+        with tarfile.open(file, "r:gz") as tar:
             file_names = tar.getnames()
-        print(file_names)
+            logger.info(file_names)
             module_dir = os.path.join(dirname, file_names[0])
             for file_name in file_names:
                 tar.extract(file_name, dirname)
 
-    return module_dir
-
+        if delete_file:
+            os.remove(file)
 
-if __name__ == "__main__":
-    # TODO(ZeyuChen) add unit test
-    link = "http://paddlehub.bj.bcebos.com/word2vec/word2vec-dim16-simple-example-1.tar.gz"
+        return module_dir
 
-    module_path = download_and_uncompress(link)
-    print("module path", module_path)
+    def download_file_and_uncompress(self,
+                                     url,
+                                     save_path=None,
+                                     save_name=None,
+                                     retry_limit=3,
+                                     delete_file=True):
+        file = self.download_file(
+            url=url,
+            save_path=save_path,
+            save_name=save_name,
+            retry_limit=retry_limit)
+        return self.uncompress(file, delete_file=delete_file)
+
+
+default_downloader = Downloader()