Browse Source

Add staging option to downloader

Kirk Trombley 2 years ago
parent
commit
0fcb860766
1 changed files with 17 additions and 11 deletions
  1. 17 11
      tools/download.py

+ 17 - 11
tools/download.py

@@ -515,8 +515,8 @@ def get_serebii_url(pkmn: Pokemon, form: Form) -> str | None:
     print(f"No Serebii URL known for {form.name}")
 
 
-async def download(session: ClientSession, url: str, filename: Path) -> tuple[str, Exception | bool]:
-  if filename.is_file():
+async def download(session: ClientSession, url: str, filename: Path, staging: bool = False) -> tuple[str, Exception | bool]:
+  if filename.is_file() or staging:
     return url, False
   try:
     async with session.get(url) as res:
@@ -528,7 +528,7 @@ async def download(session: ClientSession, url: str, filename: Path) -> tuple[st
   return url, True
 
 
-async def download_all_for_pokemon(pkmn: Pokemon, image_dir: Path) -> dict[str, dict[str, Exception | bool]]:
+async def download_all_for_pokemon(pkmn: Pokemon, image_dir: Path, staging: bool = False) -> dict[str, dict[str, Exception | bool]]:
   results = defaultdict(dict)
   async with ClientSession() as session:
     for form in pkmn.forms:
@@ -537,17 +537,18 @@ async def download_all_for_pokemon(pkmn: Pokemon, image_dir: Path) -> dict[str,
       urls.append((get_serebii_url(pkmn, form), "png"))
       # TODO more sources
       results[form.name].update(await asyncio.gather(*[
-        download(session, url, image_dir.joinpath(f"{form.name}-{i}.{ext}"))
+        download(session, url, image_dir.joinpath(
+          f"{form.name}-{i}.{ext}"), staging)
         for i, (url, ext) in enumerate(urls) if url is not None
       ]))
   return results
 
 
-async def download_all(image_dir: Path, pkmn: list[Pokemon]) -> dict[str, dict[str, Exception | bool]]:
+async def download_all(image_dir: Path, pkmn: list[Pokemon], staging: bool = False) -> dict[str, dict[str, Exception | bool]]:
   image_dir.mkdir(parents=True, exist_ok=True)
   log = {}
   for p in pkmn:
-    log.update(await download_all_for_pokemon(p, image_dir))
+    log.update(await download_all_for_pokemon(p, image_dir, staging))
   return log
 
 
@@ -787,12 +788,12 @@ KNOWN_MISSING_PNGS = ("vivillon", "furfrou", "alcremie")
 
 async def main(
   dex_file: Path, image_dir: Path, startIndex: int, endIndex: int,
-  log_skipped: bool, force_dex: bool, dex_only: bool
+  log_skipped: bool, force_dex: bool, dex_only: bool, staging: bool,
 ):
   dex = await load_pokedex(dex_file, force_dex)
   if dex_only:
     return
-  log = await download_all(image_dir, (dex[i] for i in range(startIndex, endIndex + 1)))
+  log = await download_all(image_dir, (dex[i] for i in range(startIndex, endIndex + 1)), staging)
   new_downloads = 0
   for form, result in log.items():
     for url, info in result.items():
@@ -800,8 +801,10 @@ async def main(
         if url not in KNOWN_MISSING and not (".png" in url and any(s in url for s in KNOWN_MISSING_PNGS)):
           print(f"{form}: FAILED {url} - {info}")
       elif not info:
-        if log_skipped:
-          print(f"{form}: SKIPPED {url} - {info}")
+        if staging:
+          print(url)
+        elif log_skipped:
+          print(f"{form}: SKIPPED {url}")
       else:
         print(f"{form}: SUCCESS {url}")
         new_downloads += 1
@@ -823,6 +826,9 @@ if __name__ == "__main__":
   parser.add_argument(
     "--pokedex-only", action="store_true", help="Quit before image download"
   )
+  parser.add_argument(
+    "-s", "--staging", action="store_true", help="Log URLs without downloading"
+  )
   parser.add_argument(
     "-o", "--output", default="images", type=Path, help="Image output directory"
   )
@@ -839,5 +845,5 @@ if __name__ == "__main__":
 
   asyncio.run(main(
     args.pokedex, args.output, start, end,
-    args.log_skipped, args.refresh_dex, args.pokedex_only
+    args.log_skipped, args.refresh_dex, args.pokedex_only, args.staging
   ))