Allow dumping all of Gen 1 simultaneously

pokedex/extract/rby.py

@@ -9,6 +9,7 @@ This was a pain in the ass!  Thank you SO MUCH to:
 """
 # TODO fix that docstring
 # TODO note terminology somewhere: id, index, identifier
+from collections import defaultdict
 from collections import OrderedDict
 import hashlib
 import io
@@ -1879,16 +1880,6 @@ class RBYCart:
         return Array(NUM_MOVES, PokemonCString('move_name')).parse_stream(self.stream)
 
 
-
-class RBYLoader:
-    def __init__(self, *carts):
-        self.carts = carts
-        # TODO require all the same game
-
-    def load(self):
-        pass
-
-
 # TODO would be slick to convert this to a construct...  construct
 def bitfield_to_machines(bits, machine_moves):
     machines = []
@@ -1915,22 +1906,50 @@ class WriterWrapper:
         return getattr(self.locus, key)
 
 
-def main(root):
+def main(base_root):
     # TODO does this need to take arguments?  or like, sprite mode i guess
-    carts = []
+    carts = defaultdict(dict)  # game => language => RBYCart
     for filename in sys.argv[1:]:
         cart = RBYCart(Path(filename))
-        carts.append(cart)
+        game_carts = carts[cart.game]
+        if cart.language in game_carts:
+            print(
+                "WARNING: ignoring {0.path} because it's the same game and "
+                "language ({0.game}, {0.language}) as {1.path}"
+                .format(cart, game_carts[cart.language]))
+            continue
+        game_carts[cart.language] = cart
 
-    root /= carts[0].game
+    for game, game_carts in sorted(carts.items()):
+        print()
+        print("Dumping", game)
+        if game in GAME_RELEASE_MD5SUMS:
+            got_languages = game_carts.keys()
+            expected_languages = GAME_RELEASE_MD5SUMS[game].keys()
+            extra_languages = got_languages - expected_languages
+            if extra_languages:
+                print(
+                    "WARNING: don't recognize languages {}"
+                    .format(', '.join(sorted(extra_languages))))
+
+            missing_languages = expected_languages - got_languages
+            if missing_languages:
+                print(
+                    "WARNING: missing cartridges for {} — this dump will "
+                    "be incomplete!"
+                    .format(', '.join(sorted(missing_languages))))
+
+        root = base_root / game
         root.mkdir(exist_ok=True)
 
-    #loader = RBYLoader(*carts)
+        pokemons = None
+        for language, cart in sorted(game_carts.items()):
+            if pokemons is None:
                 pokemons = OrderedDict([
                     (POKEMON_IDENTIFIERS[id + 1], schema.Pokemon())
-        for id in range(carts[0].NUM_POKEMON)
+                    for id in range(cart.NUM_POKEMON)
                 ])
-    for cart in carts:
+
             for id in range(cart.NUM_POKEMON):
                 pokemon = pokemons[POKEMON_IDENTIFIERS[id + 1]]
                 #writer = WriterWrapper(pokemon)
@@ -1960,8 +1979,8 @@ def main(root):
                 writer.base_experience = record.base_experience
                 #writer.pokedex_numbers = dict(kanto=record.pokedex_number)
 
-            # Starting moves are stored with the Pokémon; other level-up moves are
+                # Starting moves are stored with the Pokémon; other level-up
-            # stored with evolutions
+                # moves are stored with evolutions
                 level_up_moves = [
                     {1: move}
                     for move in record.initial_moveset
@@ -1980,8 +1999,9 @@ def main(root):
                     record.machines, cart.machine_moves)
 
                 # Evolution
-            # TODO alas, the species here is a number, because it's an internal id
+                # TODO alas, the species here is a number, because it's an
-            # and we switch those back using data from the game...
+                # internal id and we switch those back using data from the
+                # game...
                 evolutions = []
                 for evo_datum in cart.pokemon_evos_and_moves[id].evolutions:
                     evo = {
@@ -1993,7 +2013,9 @@ def main(root):
                     evolutions.append(evo)
                 writer.evolutions = evolutions
 
-    with (root / 'pokemon.yaml').open('w') as f:
+        fn = root / 'pokemon.yaml'
+        print('Writing', fn)
+        with fn.open('w') as f:
             f.write(Camel([schema.POKEDEX_TYPES]).dump(pokemons))