pokemon-color-search/tools/download.py

"""
Manage the logic of downloading the pokedex and source images.
"""

import re
import json
import asyncio
from dataclasses import dataclass, asdict
from collections import defaultdict

from aiohttp import ClientSession


JS_TO_JSON = re.compile(r"\b([a-zA-Z][a-zA-Z0-9]*?):")


# the dex from showdown assumes only strawberry alcremie, since
# that's what's in showdown, but we might as well add the rest
ALCREMIE_SWEETS = [
  "Strawberry", "Berry", "Love", "Star",
  "Clover", "Flower", "Ribbon",
]

# https://bulbapedia.bulbagarden.net/wiki/List_of_Pok%C3%A9mon_with_gender_differences
# there are some pokemon with notable gender diffs that the dex doesn't cover
# judgement calls made arbitrarily
GENDER_DIFFS = (
  "hippopotas", "hippowdon",
  "unfezant",
  "frillish", "jellicent",
  "pyroar",
  # meowstic, indeedee, basculegion, oinkologne are already handled in the dex
)


@dataclass
class Form:
  name: str
  traits: list[str]
  types: list[str]
  color: str


@dataclass
class Pokemon:
  num: int
  species: str
  forms: list[Form]


async def load_pokedex() -> dict:
  async with ClientSession() as session:
    async with session.get("https://play.pokemonshowdown.com/data/pokedex.js") as res:
      res.raise_for_status()
      text = await res.text("utf-8")
  # this is not json of course, but it's close
  # start by taking out the ; and definition
  cleaned = text.replace("exports.BattlePokedex = ", "").strip(";")
  # then convert the keys to strings
  converted = re.sub(JS_TO_JSON, lambda m: f'"{m.group(1)}":', cleaned)
  # and fix Type: Null
  fixed = converted.replace('""Type": Null"', '"Type: Null"')
  # then, parse it
  return json.loads(fixed)


def get_traits(species: str, kind: str) -> list[str]:
  traits = []
  if kind in ("mega", "mega-x", "mega-y", "primal"):
    traits.extend(("mega", "nostart"))
  if kind in ("gmax", "eternamax", "rapid-strike-gmax"):
    traits.extend(("gmax", "nostart"))
  if kind in ("alola", "galar", "hisui", "galar", "paldea"):
    traits.extend(("regional", kind))

  # special cases
  if species == "Tauros" and "-paldea" in kind:
    # paldean tauros has dumb names
    traits.extend(("regional", "paldea"))
  if species == "Minior" and kind != "meteor":
    # minior can only start the battle in meteor form
    traits.append("nostart")
  if species == "Darmanitan" and "zen" in kind:
    # darmanitan cannot start in zen form
    traits.append("nostart")
    if "galar" in kind:
      # also there's a galar-zen form to handle
      traits.extend(("regional", "galar"))
  if species == "Palafin" and kind == "hero":
    # palafin can only start in zero form
    traits.append("nostart")
  if species == "Gimmighoul" and kind == "roaming":
    # gimmighoul roaming is only in PGO
    traits.append("nostart")

  return list(set(traits))


def clean_dex(raw: dict) -> dict[int, Pokemon]:
  regrouped = defaultdict(list)

  for key, entry in raw.items():
    isNonstandard = entry.get("isNonstandard", None)
    baseSpecies = entry.get("baseSpecies", None)
    forme = entry.get("forme", None)
    if isNonstandard not in (None, "Past", "Unobtainable"):
      continue  # remove CAP etc.
    if baseSpecies in ("Pikachu", "Pichu") and forme is not None:
      continue  # remove pikachu spam + spiky ear pichu
    if forme == "Totem":
      continue  # remove totem pokemon

    num = entry["num"]
    # non-cosmetic forms get separate entries automatically
    # but keeping the separate unown forms would be ridiculous
    if key != "unown" and len(cosmetic := entry.get("cosmeticFormes", [])) > 0:
      cosmetic.append(f'{entry["name"]}-{entry["baseForme"]}')
      if key == "alcremie":
        # oh god this thing
        cosmetic = [
          f"{cf}-{sweet}"
          for cf in cosmetic
          for sweet in ALCREMIE_SWEETS
        ]
      regrouped[num].extend({
        **entry,
        "forme": cf.replace(" ", "-"),
        "formeKind": "cosmetic",
      } for cf in cosmetic)
    elif key in GENDER_DIFFS:
      regrouped[num].append({
        **entry,
        "forme": f'{entry["name"]}-M',
        "formeKind": "cosmetic",
      })
      regrouped[num].append({
        **entry,
        "forme": f'{entry["name"]}-F',
        "formeKind": "cosmetic",
      })
    else:
      regrouped[num].append({
        **entry,
        "forme": entry["name"],
        "formeKind": entry.get("forme", "base").lower(),
      })

  cleaned = {}
  for i in range(1, max(regrouped.keys()) + 1):
    forms = regrouped[i]
    # double check there's no skipped or duped entries
    assert len(forms) > 0 and i not in cleaned
    species = forms[0].get("baseSpecies", forms[0]["name"])
    cleaned[i] = Pokemon(
      num=i,
      species=species,
      forms=[
        Form(
          name=f.get("forme", f["name"]),
          traits=get_traits(species, f["formeKind"]),
          types=f["types"],
          color=f["color"],
        ) for f in forms
      ]
    )

  return cleaned


async def main(dex_file: str):
  # first download the pokedex
  raw_dex = await load_pokedex()
  # clean and reorganize it
  dex = clean_dex(raw_dex)
  # output dex for auditing
  with open(dex_file, "w") as out:
    json.dump({str(i): asdict(pkmn) for i, pkmn in dex.items()}, out, indent=2)
  # TODO actually progress to images


if __name__ == "__main__":
  from sys import argv
  dex_file = argv[1] if len(argv) > 1 else "data/pokedex.json"

  asyncio.run(main(dex_file))