asyncify server download

.github/workflows/linting.yml

@@ -40,7 +40,7 @@ jobs:
       - name: Install project
         run: poetry install --extras azure
 
-      - name: linting
+      - name: ruff
         run: poetry run ruff .
 
       - name: black

admin_scripts/download.py

@@ -1,10 +1,11 @@
 #!/usr/bin/env python3
 
 import argparse
+import asyncio
 import json
 from pathlib import Path
 
-import requests
+import aiohttp
 from pydantic import BaseModel, HttpUrl
 
 SERVER = "http://localhost:8000"
@@ -32,6 +33,29 @@ class Page(BaseModel):
     results: list[Level] = []
 
 
+async def download_hint(
+    session: aiohttp.ClientSession,
+    path: Path,
+    level: Level,
+    hint: Hint,
+) -> None:
+    level_dir = path / str(level.number)
+    async with session.get(hint.image.unicode_string()) as r:
+        if not r.ok:
+            print(f"Error downloading level {level.number} image {hint.number}")
+            print(await r.text())
+
+        r.raise_for_status()
+
+        assert hint.image.path is not None
+        suffix = Path(hint.image.path).suffix
+        image_file = level_dir / f"image{hint.number}{suffix}"
+
+        with image_file.open("wb") as f:
+            async for chunk, _end in r.content.iter_chunks():
+                f.write(chunk)
+
+
 def download_level(path: Path, level: Level) -> None:
     print(f"Downloading level {level.number}")
     level_dir = path / str(level.number)
@@ -51,43 +75,34 @@ def download_level(path: Path, level: Level) -> None:
         blurb_txt = level_dir / "blurb.txt"
         blurb_txt.write_text(level.description)
 
-    for hint in level.hints:
-        r = requests.get(hint.image.unicode_string(), stream=True, timeout=5)
 
-        if not r.ok:
-            print(f"Error downloading level {level.number} image {hint.number}")
-            print(r.text)
-
-        r.raise_for_status()
-
-        assert hint.image.path is not None
-        suffix = Path(hint.image.path).suffix
-        image_file = level_dir / f"image{hint.number}{suffix}"
-
-        with image_file.open("wb") as f:
-            for chunk in r.iter_content():
-                f.write(chunk)
+async def main(path: Path) -> None:
+    async with aiohttp.ClientSession() as session:
+        next_page: HttpUrl | None = HttpUrl(f"{SERVER}/api/levels")
 
+        while next_page is not None:
+            async with session.get(
+                next_page.unicode_string(),
+                auth=aiohttp.BasicAuth(USERNAME, PASSWORD),
+            ) as r:
+                if not r.ok:
+                    print("Error downloading levels")
+                    print(await r.text())
 
-def main(path: Path) -> None:
-    next_page: HttpUrl | None = HttpUrl(f"{SERVER}/api/levels")
+                r.raise_for_status()
 
-    while next_page is not None:
-        r = requests.get(
-            next_page.unicode_string(), auth=(USERNAME, PASSWORD), timeout=5
-        )
+                page = Page.model_validate_json(await r.text())
 
-        if not r.ok:
-            print("Error downloading levels")
-            print(r.text)
-
-        r.raise_for_status()
+            hint_downloads = []
+            for level in page.results:
+                download_level(path, level)
+                for hint in level.hints:
+                    hint_download = download_hint(session, path, level, hint)
+                    hint_downloads.append(hint_download)
 
-        page = Page.model_validate_json(r.text)
-        for level in page.results:
-            download_level(path, level)
+            await asyncio.gather(*hint_downloads)
 
-        next_page = page.next
+            next_page = page.next
 
 
 if __name__ == "__main__":
@@ -100,4 +115,4 @@ def main(path: Path) -> None:
     )
     args = parser.parse_args()
 
-    main(args.target)
+    asyncio.run(main(args.target))