Threadpool fixes and CLI improvements (#550)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
This commit is contained in:
Glenn Jocher 2023-01-22 17:08:08 +01:00 committed by GitHub
parent d9a0fba251
commit 21b701c4ea
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
22 changed files with 338 additions and 251 deletions

View file

@ -141,10 +141,14 @@ def download(url, dir=Path.cwd(), unzip=True, delete=True, curl=False, threads=1
dir = Path(dir)
dir.mkdir(parents=True, exist_ok=True) # make directory
if threads > 1:
pool = ThreadPool(threads)
pool.imap(lambda x: download_one(*x), zip(url, repeat(dir))) # multithreaded
pool.close()
pool.join()
# pool = ThreadPool(threads)
# pool.imap(lambda x: download_one(*x), zip(url, repeat(dir))) # multithreaded
# pool.close()
# pool.join()
with ThreadPool(threads) as pool:
pool.imap(lambda x: download_one(*x), zip(url, repeat(dir))) # multithreaded
pool.close()
pool.join()
else:
for u in [url] if isinstance(url, (str, Path)) else url:
download_one(u, dir)