Repository: DataBiosphere/getm (language: python)

Issue:
Batched downloads



Discussion / hints:

Fix patch:
diff --git a/getm/cli.py b/getm/cli.py
index bcf2a48..a68bae9 100644
--- a/getm/cli.py
+++ b/getm/cli.py
@@ -1,16 +1,69 @@
-import os
+"""Download data from http(s) URLs to the local file system, verifying integrity."""
 import sys
+import json
+import logging
+import argparse
+from math import ceil
+from concurrent.futures import Future, ProcessPoolExecutor, as_completed
+from typing import List, Optional
 
-from getm import iter_content
+from getm import urlopen, iter_content, default_chunk_size
 from getm.http import http
 from getm.utils import checksum_for_url, indirect_open
-from getm.progress import ProgressBar
+from getm.progress import ProgressIndicator, ProgressBar, ProgressLogger
 
 
-def do_download(url, name, filepath):
+MB = 1024 * 1024
+logging.basicConfig()
+logging.getLogger().setLevel(logging.INFO)
+
+class Progress:
+    progress_class: type = ProgressBar
+
+    @classmethod
+    def get(cls, name: str, url: str):
+        sz = http.size(url)
+        if cls.progress_class == ProgressBar:
+            incriments = 40
+        else:
+            incriments = ceil(sz / default_chunk_size / 2)
+        return cls.progress_class(name, sz, incriments)
+
+def download(url_info: List[dict],
+             oneshot_concurrency: int=4,
+             multipart_concurrency: int=2,
+             multipart_threshold=default_chunk_size):
+    assert 1 <= oneshot_concurrency
+    assert 1 <= multipart_concurrency
+    with ProcessPoolExecutor(max_workers=oneshot_concurrency) as oneshot_executor:
+        with ProcessPoolExecutor(max_workers=multipart_concurrency) as multipart_executor:
+            for info in url_info:
+                url = info['url']
+                filepath = info.get('filepath') or http.name(url)
+                if multipart_threshold >= http.size(url):
+                    oneshot_executor.submit(oneshot, url, filepath)
+                else:
+                    multipart_executor.submit(multipart, url, filepath)
+
+def oneshot(url: str, filepath: str):
+    expected_cs, cs = checksum_for_url(url)
+    assert expected_cs and cs, "No checksum information available!"
+    with urlopen(url, concurrency=None) as handle:
+        data = handle.read()
+        try:
+            cs.update(data)
+            assert cs.matches(expected_cs), "Checksum failed!"
+            with open(filepath, "wb", buffering=0) as fh:
+                fh.write(data)
+            with Progress.get(filepath, url) as progress:
+                progress.add(len(data))
+        finally:
+            data.release()
+
+def multipart(url: str, filepath: str):
     expected_cs, cs = checksum_for_url(url)
-    assert expected_cs is not None, "No checksum information available!"
-    with ProgressBar(name, http.size(url)) as progress:
+    assert expected_cs and cs, "No checksum information available!"
+    with Progress.get(filepath, url) as progress:
         with indirect_open(filepath) as handle:
             for part in iter_content(url, concurrency=1):
                 handle.write(part)
@@ -18,8 +71,46 @@ def do_download(url, name, filepath):
                 progress.add(len(part))
             assert cs.matches(expected_cs), "Checksum failed!"
 
+manifest_arg_doc = """Download URls as specified in a local json FILE
+The json is expected to be an array of elements of the following structure. The only required field is 'url'.
+{
+  'url':
+  'filepath': # local name on file system
+}
+"""
+
 def main():
-    url = sys.argv[1]
-    name = sys.argv[2]
-    filepath = os.path.realpath(sys.argv[2])
-    do_download(url, name, filepath)
+    """This is the main entry point for the CLI."""
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
+    parser.add_argument("url",
+                        nargs="?",
+                        metavar="url",
+                        type=str,
+                        default=None,
+                        help="URL to download. This argument has greatest precedence.")
+    parser.add_argument("--filepath", "-O", help="Local file path")
+    parser.add_argument("--manifest", "-m", "-i", help=manifest_arg_doc)
+    parser.add_argument("--oneshot-concurrency", default=4)
+    parser.add_argument("--multipart-concurrency", default=2)
+    parser.add_argument("--multipart-threshold", default=default_chunk_size)
+    args = parser.parse_args()
+
+    if not (args.url or args.manifest):
+        parser.print_usage()
+        sys.exit(1)
+
+    if args.url:
+        info = dict(url=args.url)
+        if args.filepath:
+            info['filepath'] = args.filepath
+        url_info = [info]
+    else:
+        with open(args.manifest) as fh:
+            url_info = json.loads(fh.read())
+
+    if 1 == len(url_info):
+        Progress.progress_class = ProgressBar
+    else:
+        Progress.progress_class = ProgressLogger
+
+    download(url_info, args.oneshot_concurrency, args.multipart_concurrency, args.multipart_threshold)


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
new file mode 100644
index 0000000..e1d03ef
--- /dev/null
+++ b/tests/test_cli.py
@@ -0,0 +1,126 @@
+#!/usr/bin/env python
+import os
+import sys
+import unittest
+import contextlib
+from uuid import uuid4
+from hashlib import md5
+from unittest import mock
+from tempfile import TemporaryDirectory
+
+pkg_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))  # noqa
+sys.path.insert(0, pkg_root)  # noqa
+
+from getm.cli import download, oneshot, multipart
+from getm.checksum import MD5
+
+from tests.infra import suppress_warnings
+from tests.infra.server import ThreadedLocalServer, SilentHandler
+
+
+class Server:
+    server: ThreadedLocalServer
+    host: str
+    data = dict()
+
+    @classmethod
+    def set_data(cls, length: int):
+        path = f"/{uuid4()}"
+        cls.data[path] = os.urandom(length)
+        return cls.host + path, cls.data[path]
+
+class MockExecutor:
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    def submit(self, func, *args, **kwargs):
+        func(*args, **kwargs)
+
+def setUpModule():
+    class Handler(SilentHandler):
+        def do_HEAD(self, *args, **kwargs):
+            data = Server.data[self.path]
+            self.send_response(200)
+            self.send_header("Content-Length", len(data))
+            self.send_header("Content-MD5", md5(data).hexdigest())
+            self.end_headers()
+
+        def do_GET(self, *args, **kwargs):
+            self.do_HEAD(*args, **kwargs)
+            self.wfile.write(Server.data[self.path])
+
+    Server.server = ThreadedLocalServer(Handler)
+    Server.server.start()
+    Server.host = Server.server.host
+
+def tearDownModule():
+    Server.server.shutdown()
+
+@mock.patch("getm.cli.Progress")
+class TestCLI(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.temp_dir = TemporaryDirectory()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.temp_dir.cleanup()
+
+    def setUp(self):
+        suppress_warnings()
+        self.filepath = f"{self.temp_dir.name}/{uuid4()}"
+        Server.data = dict()
+
+    def test_download(self, *args):
+        multipart_threshold = 7
+        oneshot_sizes = 3 * [1]
+        multipart_sizes = 5 * [1 + multipart_threshold]
+        url_info = [dict(url=Server.set_data(size)[0], filepath=f"{self.temp_dir.name}/{uuid4()}")
+                    for size in oneshot_sizes + multipart_sizes]
+
+        with self.subTest("routing"):
+            with mock.patch("getm.cli.ProcessPoolExecutor", MockExecutor):
+                with mock.patch("getm.cli.oneshot") as mock_oneshot:
+                    with mock.patch("getm.cli.multipart") as mock_multipart:
+                        download(url_info, multipart_threshold=multipart_threshold)
+                        self.assertEqual(len(oneshot_sizes), len(mock_oneshot.call_args_list))
+                        self.assertEqual(len(multipart_sizes), len(mock_multipart.call_args_list))
+
+        for oneshot_concurrency, multipart_concurrency in [(0, 0), (1, 0), (0, 1)]:
+            with self.subTest("assertions",
+                              oneshot_concurrency=oneshot_concurrency,
+                              multipart_concurrency=multipart_concurrency):
+                with self.assertRaises(AssertionError):
+                    download(url_info, oneshot_concurrency, multipart_concurrency)
+
+        with self.subTest("download"):
+            download(url_info, multipart_threshold=multipart_threshold)
+            for info in url_info:
+                url = info['url']
+                path = "/" + info['url'].rsplit("/", 1)[-1]
+                with open(info['filepath'], "rb") as fh:
+                    self.assertEqual(Server.data[path], fh.read())
+
+    def test_oneshot(self, *args):
+        url, expected_data = Server.set_data(1021)
+        expected_cs, cs = md5(expected_data).hexdigest(), MD5()
+        oneshot(url, self.filepath)
+        with open(self.filepath, "rb") as fh:
+            self.assertEqual(expected_data, fh.read())
+
+    @mock.patch("getm.default_chunk_size", 1021)
+    def test_multipart(self, *args):
+        url, expected_data = Server.set_data(999983)
+        expected_cs, cs = md5(expected_data).hexdigest(), MD5()
+        multipart(url, self.filepath)
+        with open(self.filepath, "rb") as fh:
+            self.assertEqual(expected_data, fh.read())
+
+if __name__ == '__main__':
+    unittest.main()