|
@@ -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
|
|
|
+ ))
|