From d24f92a3faee80cd11075be697697ce5f190c0fd Mon Sep 17 00:00:00 2001
From: sheepish <sheepish@orpheus.network>
Date: Wed, 3 May 2023 20:35:27 +0000
Subject: [PATCH] add populate_cache.py

---
 populate_cache.py | 74 +++++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt  |  3 +-
 2 files changed, 76 insertions(+), 1 deletion(-)
 create mode 100644 populate_cache.py

diff --git a/populate_cache.py b/populate_cache.py
new file mode 100644
index 0000000..39f8f19
--- /dev/null
+++ b/populate_cache.py
@@ -0,0 +1,74 @@
+import asyncio
+import argparse
+import aiofiles
+from aiofiles.os import path
+from base64 import urlsafe_b64encode
+import tqdm
+from os.path import join as path_join
+
+from main import create_filename, requester
+
+
+async def download(base_dir: str, url: str):
+    img_url = urlsafe_b64encode(url.encode('ascii'))
+    file_path = path_join(base_dir, create_filename(img_url).decode('ascii'))
+    if await path.isfile(file_path):
+        return
+    await requester.save(img_url.decode('ascii'), file_path)
+
+
+async def main():
+    parser = argparse.ArgumentParser(description='Fetch images from text file')
+    parser.add_argument('dest', help='base path to store images at')
+    parser.add_argument('file', help='text file with links to import, one per line')
+    parser.add_argument('--disable-parallel', '-x', action='store_true',
+                        help='do not do multiple requests at the same time')
+    args = parser.parse_args()
+
+    if args.disable_parallel:
+        worker_slots = 1
+    else:
+        worker_slots = 100
+    sem = asyncio.Semaphore(worker_slots)
+    success = []
+    failed = []
+    pbar = tqdm.tqdm(unit='req', postfix={'success': 0, 'failed': 0})
+
+    def callback(task):
+        try:
+            task.result()
+        except Exception:
+            failed.append(task.get_name())
+        else:
+            success.append(task.get_name())
+        finally:
+            sem.release()
+            pbar.update(1)
+            pbar.set_postfix(success=len(success), failed=len(failed))
+
+    i = 0
+    async with aiofiles.open(args.file, 'r') as fh:
+        async for line in fh:
+            i += 1
+            await sem.acquire()
+            url = line.rstrip('\n')
+            task = asyncio.create_task(download(args.dest, url))
+            task.set_name(url)
+            task.add_done_callback(callback)
+
+    for x in range(worker_slots):
+        # acquire all slots so no more tasks can be active
+        await sem.acquire()
+
+    pbar.close()
+
+    async with aiofiles.open(args.file + '.success', 'w') as fh:
+        await fh.writelines(x + '\n' for x in success)
+    async with aiofiles.open(args.file + '.failed', 'w') as fh:
+        await fh.writelines(x + '\n' for x in failed)
+
+    print(f"all done. total {i}, success: {len(success)}, failed: {len(failed)}")
+
+
+if __name__ == '__main__':
+    asyncio.run(main())
diff --git a/requirements.txt b/requirements.txt
index 287b57e..f9b3276 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,5 @@
 blacksheep~=1.2.14
 uvicorn~=0.22.0
 aiofiles~=23.1.0
-httpx~=0.24.0
\ No newline at end of file
+httpx~=0.24.0
+tqdm~=4.65.0
\ No newline at end of file