diff --git a/osfclient/__main__.py b/osfclient/__main__.py index 3e277ad..34754c2 100644 --- a/osfclient/__main__.py +++ b/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='?') diff --git a/osfclient/cli.py b/osfclient/cli.py index 6260fe0..8df24b1 100644 --- a/osfclient/cli.py +++ b/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. diff --git a/osfclient/models/file.py b/osfclient/models/file.py index 6ca98cd..0eda515 100644 --- a/osfclient/models/file.py +++ b/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 " diff --git a/requirements.txt b/requirements.txt index f229360..5bb8c66 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ requests +tqdm diff --git a/setup.py b/setup.py index 37eedbd..7faf609 100644 --- a/setup.py +++ b/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