Browse Source

add arg controls to download

Kirk Trombley 2 years ago
parent
commit
bd89b6789f
1 changed files with 48 additions and 21 deletions
  1. 48 21
      tools/download.py

+ 48 - 21
tools/download.py

@@ -168,8 +168,8 @@ def clean_dex(raw: dict) -> dict[int, Pokemon]:
   }
 
 
-async def load_pokedex(dex_file: str) -> dict:
-  if Path(dex_file).is_file():
+async def load_pokedex(dex_file: Path, force_dex: bool) -> dict:
+  if dex_file.is_file() and not force_dex:
     with open(dex_file) as infile:
       loaded = json.load(infile)
     dex = {
@@ -202,7 +202,7 @@ SHOWDOWN_REPLACEMENTS = [
   ("nidoran-m", "nidoranm"),  # nidoran is a special case
   ("-f", "f"),  # gender diff forms
   (re.compile(r"-m$"), ""),  # gender diff forms
-  (re.compile(r"^ho-oh$"), "hooh"), # Ho-oh special case
+  (re.compile(r"^ho-oh$"), "hooh"),  # Ho-oh special case
 ]
 
 
@@ -221,8 +221,8 @@ def get_showdown_urls(form: Form) -> list[tuple[str, str]]:
   ]
 
 
-async def download(session: ClientSession, url: str, filename: str) -> tuple[str, Exception | bool]:
-  if Path(filename).is_file():
+async def download(session: ClientSession, url: str, filename: Path) -> tuple[str, Exception | bool]:
+  if filename.is_file():
     return url, False
   try:
     async with session.get(url) as res:
@@ -234,7 +234,7 @@ async def download(session: ClientSession, url: str, filename: str) -> tuple[str
   return url, True
 
 
-async def download_all_for_pokemon(pkmn: Pokemon, image_dir: str) -> dict[str, dict[str, Exception | bool]]:
+async def download_all_for_pokemon(pkmn: Pokemon, image_dir: Path) -> dict[str, dict[str, Exception | bool]]:
   results = defaultdict(dict)
   async with ClientSession() as session:
     for form in pkmn.forms:
@@ -242,22 +242,27 @@ async def download_all_for_pokemon(pkmn: Pokemon, image_dir: str) -> dict[str, d
       urls += get_showdown_urls(form)
       # TODO more sources
       results[form.name].update(await asyncio.gather(*[
-        download(session, url, f"{image_dir}/{form.name}-{i}.{ext}")
+        download(session, url, image_dir.joinpath(f"{form.name}-{i}.{ext}"))
         for i, (url, ext) in enumerate(urls)
       ]))
   return results
 
 
-async def download_all(image_dir: str, pkmn: list[Pokemon]) -> dict[str, dict[str, Exception | bool]]:
-  Path(image_dir).mkdir(parents=True, exist_ok=True)
+async def download_all(image_dir: Path, pkmn: list[Pokemon]) -> 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))
   return log
 
 
-async def main(dex_file: str, image_dir: str, startIndex: int, endIndex: int, log_skipped: bool):
-  dex = await load_pokedex(dex_file)
+async def main(
+  dex_file: Path, image_dir: Path, startIndex: int, endIndex: int,
+  log_skipped: bool, force_dex: bool, dex_only: 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)))
   new_downloads = 0
   for form, result in log.items():
@@ -273,13 +278,35 @@ async def main(dex_file: str, image_dir: str, startIndex: int, endIndex: int, lo
 
 
 if __name__ == "__main__":
-  from sys import argv
-  # TODO make this an arg parser
-  dex_file = argv[1] if len(argv) > 1 else "data/pokedex.json"
-  image_dir = argv[2] if len(argv) > 2 else "images"
-  start, end = map(int, (
-    argv[3] if len(argv) > 3 else "1-151"
-  ).split("-")[0:2])
-  log_skipped = len(argv) > 4 and argv[4].lower() == 'true'
-
-  asyncio.run(main(dex_file, image_dir, start, end, log_skipped))
+  from argparse import ArgumentParser
+  parser = ArgumentParser(
+    prog="Image Retriever",
+    description="Retrieve pokedex and images",
+  )
+  parser.add_argument(
+    "-d", "--pokedex", default="data/pokedex.json", type=Path, help="Pokedex file"
+  )
+  parser.add_argument(
+    "--refresh-dex", action="store_true", help="Update the pokedex"
+  )
+  parser.add_argument(
+    "--pokedex-only", action="store_true", help="Quit before image download"
+  )
+  parser.add_argument(
+    "-o", "--output", default="images", type=Path, help="Image output directory"
+  )
+  parser.add_argument(
+    "--log-skipped", action="store_true", help="Log skipped images"
+  )
+  parser.add_argument(
+    "bounds", type=lambda a: map(int, a.split("-")), default="1-151", nargs="?",
+    help="Range of dex numbers to download, inclusive"
+  )
+
+  args = parser.parse_args()
+  start, end = args.bounds
+
+  asyncio.run(main(
+    args.pokedex, args.output, start, end,
+    args.log_skipped, args.refresh_dex, args.pokedex_only
+  ))