Fix streaming bug osfclient#55 and add progres bar

osfclient/__main__.py

@@ -38,7 +38,6 @@ def main():
         'clone', description=clone.__doc__,
         formatter_class=argparse.RawDescriptionHelpFormatter)
     clone_parser.set_defaults(func=clone)
-    clone_parser.add_argument('project', help='OSF project ID')
     clone_parser.add_argument('output', help='Write files to this directory',
                               default=None, nargs='?')
 

osfclient/cli.py

@@ -8,6 +8,8 @@
 except ImportError:
     import ConfigParser as configparser
 
+from tqdm import tqdm
+
 from .api import OSF
 from .utils import norm_remote_path, split_storage
 
@@ -79,20 +81,23 @@ def clone(args):
     if args.output is not None:
         output_dir = args.output
 
-    for store in project.storages:
-        prefix = os.path.join(output_dir, store.name)
+    with tqdm(unit='files') as pbar:
+        for store in project.storages:
+            prefix = os.path.join(output_dir, store.name)
 
-        for file_ in store.files:
-            path = file_.path
-            if path.startswith('/'):
-                path = path[1:]
+            for file_ in store.files:
+                path = file_.path
+                if path.startswith('/'):
+                    path = path[1:]
 
-            path = os.path.join(prefix, path)
-            directory, _ = os.path.split(path)
-            os.makedirs(directory, exist_ok=True)
+                path = os.path.join(prefix, path)
+                directory, _ = os.path.split(path)
+                os.makedirs(directory, exist_ok=True)
 
-            with open(path, "wb") as f:
-                file_.write_to(f)
+                with open(path, "wb") as f:
+                    file_.write_to(f)
+
+                pbar.update()
 
 
 def fetch(args):
@@ -128,6 +133,9 @@ def fetch(args):
             with open(local_path, 'wb') as fp:
                 file_.write_to(fp)
 
+            # only fetching one file so we are done
+            break
+
 
 def list_(args):
     """List all files from all storages for project.

osfclient/models/file.py

@@ -1,9 +1,23 @@
-import shutil
+from tqdm import tqdm
 
 from .core import OSFCore
 from ..exceptions import FolderExistsException
 
 
+def copyfileobj(fsrc, fdst, total, length=16*1024):
+    """Copy data from file-like object fsrc to file-like object fdst
+
+    This is like shutil.copyfileobj but with a progressbar.
+    """
+    with tqdm(unit='bytes', total=total, unit_scale=True) as pbar:
+        while 1:
+            buf = fsrc.read(length)
+            if not buf:
+                break
+            fdst.write(buf)
+            pbar.update(len(buf))
+
+
 class File(OSFCore):
     def _update_attributes(self, file):
         if not file:
@@ -35,10 +49,11 @@ def write_to(self, fp):
         if 'b' not in fp.mode:
             raise ValueError("File has to be opened in binary mode.")
 
-        response = self._get(self._download_url)
+        response = self._get(self._download_url, stream=True)
         if response.status_code == 200:
             response.raw.decode_content = True
-            shutil.copyfileobj(response.raw, fp)
+            copyfileobj(response.raw, fp,
+                        int(response.headers['Content-Length']))
 
         else:
             raise RuntimeError("Response has status "

requirements.txt

@@ -1 +1,2 @@
 requests
+tqdm

setup.py

@@ -42,7 +42,7 @@
     # your project is installed. For an analysis of "install_requires" vs pip's
     # requirements files see:
     # https://packaging.python.org/en/latest/requirements.html
-    install_requires=['requests'],
+    install_requires=['requests', 'tqdm'],
 
     # To provide executable scripts, use entry points in preference to the
     # "scripts" keyword. Entry points provide cross-platform support and allow