"""Output data models matching the existing seed JSON format.""" from __future__ import annotations from dataclasses import dataclass, field @dataclass class Encounter: pokeapi_id: int pokemon_name: str method: str encounter_rate: int min_level: int max_level: int def to_dict(self) -> dict: return { "pokeapi_id": self.pokeapi_id, "pokemon_name": self.pokemon_name, "method": self.method, "encounter_rate": self.encounter_rate, "min_level": self.min_level, "max_level": self.max_level, } @dataclass class Route: name: str order: int encounters: list[Encounter] = field(default_factory=list) children: list[Route] = field(default_factory=list) def to_dict(self) -> dict: d: dict = { "name": self.name, "order": self.order, "encounters": [e.to_dict() for e in self.encounters], } if self.children: d["children"] = [c.to_dict() for c in self.children] return d @dataclass class Game: name: str slug: str generation: int region: str release_year: int color: str | None = None def to_dict(self) -> dict: return { "name": self.name, "slug": self.slug, "generation": self.generation, "region": self.region, "release_year": self.release_year, "color": self.color, } @dataclass class Pokemon: pokeapi_id: int national_dex: int name: str types: list[str] sprite_url: str def to_dict(self) -> dict: return { "pokeapi_id": self.pokeapi_id, "national_dex": self.national_dex, "name": self.name, "types": self.types, "sprite_url": self.sprite_url, }