Files
nuzlocke-tracker/tools/import-pokedb/import_pokedb/output.py

239 lines
7.5 KiB
Python
Raw Normal View History

"""Output seed JSON files in the existing format."""
from __future__ import annotations
import json
import sys
from pathlib import Path
from typing import Any
from .loader import SeedConfig
from .mappings import PokemonMapper
from .models import Encounter, Route
from .sprites import sprite_path_for_pokemon
# ---------------------------------------------------------------------------
# Route ordering
# ---------------------------------------------------------------------------
def _get_route_order(config: SeedConfig, vg_key: str) -> list[str] | None:
"""Get the route order list for a version group, resolving aliases."""
return config.route_order.get(vg_key)
def _route_sort_key(name: str, order_list: list[str]) -> tuple[int, str]:
"""Compute sort key for a route name against an ordered list."""
for i, ordered_name in enumerate(order_list):
if name == ordered_name or name.startswith(ordered_name + " ("):
return (i, name)
return (len(order_list), name)
def sort_routes(routes: list[Route], config: SeedConfig, vg_key: str) -> list[Route]:
"""Sort routes by game progression order and assign sequential order values."""
order_list = _get_route_order(config, vg_key)
if order_list:
routes.sort(key=lambda r: _route_sort_key(r.name, order_list))
else:
print(f" Warning: No route order for {vg_key}, using default order", file=sys.stderr)
# Assign sequential order values
order = 1
for route in routes:
route.order = order
order += 1
for child in route.children:
child.order = order
order += 1
return routes
# ---------------------------------------------------------------------------
# Special encounters
# ---------------------------------------------------------------------------
def _get_special_encounters(
config: SeedConfig,
vg_key: str,
) -> dict[str, list[dict[str, Any]]] | None:
"""Get special encounters for a version group, resolving aliases."""
if not config.special_encounters:
return None
encounters = config.special_encounters.get("encounters", {})
aliases = config.special_encounters.get("aliases", {})
if vg_key in encounters:
return encounters[vg_key]
if vg_key in aliases:
alias_target = aliases[vg_key]
if alias_target in encounters:
return encounters[alias_target]
return None
def merge_special_encounters(
routes: list[Route],
config: SeedConfig,
vg_key: str,
pokemon_mapper: PokemonMapper | None = None,
) -> list[Route]:
"""Merge special encounters (starters, gifts, fossils) into routes."""
special_data = _get_special_encounters(config, vg_key)
if not special_data:
return routes
# Build lookup: route name → route (including children)
route_map: dict[str, Route] = {}
for route in routes:
route_map[route.name] = route
for child in route.children:
route_map[child.name] = child
for location_name, encounter_dicts in special_data.items():
encounters = []
for e in encounter_dicts:
# Look up proper display name from pokemon mapper if available
pokemon_name = e["pokemon_name"]
if pokemon_mapper:
info = pokemon_mapper.lookup_by_id(e["pokeapi_id"])
if info:
pokemon_name = info[1]
encounters.append(Encounter(
pokeapi_id=e["pokeapi_id"],
pokemon_name=pokemon_name,
method=e["method"],
encounter_rate=e["encounter_rate"],
min_level=e["min_level"],
max_level=e["max_level"],
))
if location_name in route_map:
route_map[location_name].encounters.extend(encounters)
else:
new_route = Route(name=location_name, order=0, encounters=encounters)
routes.append(new_route)
route_map[location_name] = new_route
return routes
# ---------------------------------------------------------------------------
# JSON output
# ---------------------------------------------------------------------------
def write_json(path: Path, data: Any) -> None:
"""Write data as formatted JSON with trailing newline."""
content = json.dumps(data, indent=2, ensure_ascii=False)
path.write_text(content + "\n", encoding="utf-8")
print(f" -> {path}")
def write_game_json(routes: list[Route], output_dir: Path, game_slug: str) -> None:
"""Write a per-game route/encounter JSON file."""
data = [r.to_dict() for r in routes]
write_json(output_dir / f"{game_slug}.json", data)
def write_games_json(config: SeedConfig, output_dir: Path) -> None:
"""Write games.json from version_groups config."""
games = []
for vg_info in config.version_groups.values():
for game_info in vg_info.get("games", {}).values():
games.append({
"name": game_info["name"],
"slug": game_info["slug"],
"generation": vg_info["generation"],
"region": vg_info["region"],
"release_year": game_info["release_year"],
"color": game_info.get("color"),
})
write_json(output_dir / "games.json", games)
print(f" Wrote {len(games)} games")
def write_pokemon_json(
pokemon_mapper: PokemonMapper,
encountered_form_ids: set[str],
sprite_map: dict[str, str],
output_dir: Path,
) -> None:
"""Write pokemon.json with all pokemon referenced in encounters."""
seen_ids: set[int] = set()
pokemon_list: list[dict[str, Any]] = []
for form_id in sorted(encountered_form_ids):
info = pokemon_mapper.lookup(form_id)
if info is None:
continue
pokeapi_id, name = info
if pokeapi_id in seen_ids:
continue
seen_ids.add(pokeapi_id)
# Get additional data from PokeDB form record
form_data = pokemon_mapper.get_form_data(form_id)
national_dex = form_data.get("ndex_id", pokeapi_id) if form_data else pokeapi_id
# Types from PokeDB form data
types = _extract_types(form_data) if form_data else []
# Sprite URL
sprite_url: str | None = None
if form_id in sprite_map:
sprite_url = sprite_path_for_pokemon(pokeapi_id)
pokemon_list.append({
"pokeapi_id": pokeapi_id,
"national_dex": national_dex,
"name": name,
"types": types,
"sprite_url": sprite_url,
})
# Sort by pokeapi_id
pokemon_list.sort(key=lambda p: p["pokeapi_id"])
write_json(output_dir / "pokemon.json", pokemon_list)
print(f" Wrote {len(pokemon_list)} pokemon")
# PokeDB type IDs → type names (from PokeDB's type system)
_TYPE_ID_MAP: dict[int, str] = {
1: "normal",
2: "fighting",
3: "flying",
4: "poison",
5: "ground",
6: "rock",
7: "bug",
8: "ghost",
9: "steel",
10: "fire",
11: "water",
12: "grass",
13: "electric",
14: "psychic",
15: "ice",
16: "dragon",
17: "dark",
18: "fairy",
}
def _extract_types(form_data: dict[str, Any]) -> list[str]:
"""Extract type names from a PokeDB form record."""
types = []
type1_id = form_data.get("type_1_id")
type2_id = form_data.get("type_2_id")
if type1_id and type1_id in _TYPE_ID_MAP:
types.append(_TYPE_ID_MAP[type1_id])
if type2_id and type2_id in _TYPE_ID_MAP:
types.append(_TYPE_ID_MAP[type2_id])
return types