Update struct

- Update for Generation 5. This means SaveFilePokemon becomes an abstract
  base, with SaveFilePokemonGen4 and SaveFilePokemonGen5 subclasses.
  The pokemon_struct is changed to a structure factory.
  Diff best viewed with the ignore whitespace setting.
- Allow creating a SaveFilePokemon with a zeroed-out blob (to be filled
  in later). Several fields are zeroed out in gen 5 so this helps the
  above.

pokedex/struct/__init__.py

@@ -13,7 +13,7 @@ import struct
 from pokedex.db import tables
 from pokedex.formulae import calculated_hp, calculated_stat
 from pokedex.compatibility import namedtuple, permutations
-from pokedex.struct._pokemon_struct import pokemon_struct
+from pokedex.struct._pokemon_struct import make_pokemon_struct
 
 def pokemon_prng(seed):
     u"""Creates a generator that simulates the main Pokémon PRNG."""
@@ -22,26 +22,32 @@ def pokemon_prng(seed):
         seed &= 0xFFFFFFFF
         yield seed >> 16
 
-
 class SaveFilePokemon(object):
-    u"""Represents an individual Pokémon, from the game's point of view.
+    u"""Base class for an individual Pokémon, from the game's point of view.
 
     Handles translating between the on-disk encrypted form, the in-RAM blob
     (also used by pokesav), and something vaguely intelligible.
     """
-
     Stat = namedtuple('Stat', ['stat', 'base', 'gene', 'exp', 'calc'])
 
-    def __init__(self, blob, encrypted=False):
+    def __init__(self, blob=None, encrypted=False, session=None):
         u"""Wraps a Pokémon save struct in a friendly object.
 
         If `encrypted` is True, the blob will be decrypted as though it were an
         on-disk save.  Otherwise, the blob is taken to be already decrypted and
         is left alone.
 
-        `session` is an optional database session.
+        `session` is an optional database session. Either give it or fill it
+            later with `use_database_session`
         """
 
+        try:
+            self.generation_id
+        except AttributeError:
+            raise NotImplementedError(
+                "Use generation-specific subclass of SaveFilePokemon")
+
+        if blob:
             if encrypted:
                 # Decrypt it.
                 # Interpret as one word (pid), followed by a bunch of shorts
@@ -58,7 +64,12 @@ class SaveFilePokemon(object):
                 # Already decrypted
                 self.blob = blob
 
-        self.structure = pokemon_struct.parse(self.blob)
+            self.structure = self.pokemon_struct.parse(self.blob)
+        else:
+            self.structure = self.pokemon_struct.parse('\0' * (32 * 4 + 8))
+
+        if session:
+            self.use_database_session(session)
 
     @property
     def as_struct(self):
@@ -98,20 +109,25 @@ class SaveFilePokemon(object):
 
     def use_database_session(self, session):
         """Remembers the given database session, and prefetches a bunch of
-        database stuff.  Gotta call this before you use the database properties
+        database stuff.  Gotta call this (or give it to `__init__`) before
-        like `species`, etc.
+        you use the database properties like `species`, etc.
         """
         self._session = session
 
         st = self.structure
+
+        if st.national_id:
             self._pokemon = session.query(tables.Pokemon).get(st.national_id)
             self._pokemon_form = session.query(tables.PokemonForm) \
                 .with_parent(self._pokemon) \
-            .filter_by(name=st.alternate_form) \
+                .filter_by(form_identifier=st.alternate_form) \
                 .one()
+        else:
+            self._pokemon = self._pokemon_form = None
         self._ability = self._session.query(tables.Ability).get(st.ability_id)
 
-        growth_rate = self._pokemon.evolution_chain.growth_rate
+        if self._pokemon:
+            growth_rate = self._pokemon.species.growth_rate
             self._experience_rung = session.query(tables.Experience) \
                 .filter(tables.Experience.growth_rate == growth_rate) \
                 .filter(tables.Experience.experience <= st.exp) \
@@ -131,6 +147,7 @@ class SaveFilePokemon(object):
             self._held_item = session.query(tables.ItemGameIndex) \
                 .filter_by(generation_id = 4, game_index = st.held_item_id).one().item
 
+        if self._pokemon:
             self._stats = []
             for pokemon_stat in self._pokemon.stats:
                 structure_name = pokemon_stat.stat.name.lower().replace(' ', '_')
@@ -156,6 +173,8 @@ class SaveFilePokemon(object):
                 )
 
                 self._stats.append(stat_tup)
+        else:
+            self._stats = [0] * 6
 
 
         move_ids = (
@@ -171,8 +190,11 @@ class SaveFilePokemon(object):
 
         if st.hgss_pokeball >= 17:
             pokeball_id = st.hgss_pokeball - 17 + 492
-        else:
+        elif st.dppt_pokeball:
             pokeball_id = st.dppt_pokeball
+        else:
+            pokeball_id = None
+        if pokeball_id:
             self._pokeball = session.query(tables.ItemGameIndex) \
                 .filter_by(generation_id = 4, game_index = pokeball_id).one().item
 
@@ -182,15 +204,25 @@ class SaveFilePokemon(object):
         self._egg_location = None
         if egg_loc_id:
             self._egg_location = session.query(tables.LocationGameIndex) \
-                .filter_by(generation_id = 4, game_index = egg_loc_id).one().location
+                .filter_by(generation_id = self.generation_id, game_index = egg_loc_id).one().location
 
+        if met_loc_id:
             self._met_location = session.query(tables.LocationGameIndex) \
-            .filter_by(generation_id = 4, game_index = met_loc_id).one().location
+                .filter_by(generation_id = self.generation_id, game_index = met_loc_id).one().location
+        else:
+            self._met_location = None
 
     @property
     def species(self):
-        # XXX forme!
+        return self._pokemon_form.species
-        return self._pokemon
+
+    @property
+    def pokemon(self):
+        return self._pokemon_form.pokemon
+
+    @property
+    def form(self):
+        return self._pokemon_form
 
     @property
     def species_form(self):
@@ -208,16 +240,6 @@ class SaveFilePokemon(object):
     def met_location(self):
         return self._met_location
 
-    @property
-    def shiny_leaves(self):
-        return (
-            self.structure.shining_leaves.leaf1,
-            self.structure.shining_leaves.leaf2,
-            self.structure.shining_leaves.leaf3,
-            self.structure.shining_leaves.leaf4,
-            self.structure.shining_leaves.leaf5,
-        )
-
     @property
     def level(self):
         return self._experience_rung.level
@@ -319,3 +341,42 @@ class SaveFilePokemon(object):
                 words[i] ^= next(prng)
 
         return
+
+
+class SaveFilePokemonGen4(SaveFilePokemon):
+    generation_id = 4
+    pokemon_struct = make_pokemon_struct(generation=generation_id)
+
+    @property
+    def shiny_leaves(self):
+        return (
+            self.structure.shining_leaves.leaf1,
+            self.structure.shining_leaves.leaf2,
+            self.structure.shining_leaves.leaf3,
+            self.structure.shining_leaves.leaf4,
+            self.structure.shining_leaves.leaf5,
+        )
+
+
+class SaveFilePokemonGen5(SaveFilePokemon):
+    generation_id = 5
+    pokemon_struct = make_pokemon_struct(generation=generation_id)
+
+    def use_database_session(self, session):
+        super(SaveFilePokemonGen5, self).use_database_session(session)
+
+        st = self.structure
+
+        if st.nature_id:
+            self._nature = session.query(tables.Nature) \
+                .filter_by(game_index = st.nature_id).one()
+
+    @property
+    def nature(self):
+        return self._nature
+
+
+save_file_pokemon_classes = {
+    4: SaveFilePokemonGen4,
+    5: SaveFilePokemonGen5,
+}

pokedex/struct/_pokemon_struct.py

@@ -578,10 +578,32 @@ class PokemonFormAdapter(Adapter):
         return forms.index(obj) << 3
 
 
-
-# And here we go.
 # Docs: http://projectpokemon.org/wiki/Pokemon_NDS_Structure
-pokemon_struct = Struct('pokemon_struct',
+# http://projectpokemon.org/wiki/Pokemon_Black/White_NDS_Structure
+# http://projectpokemon.org/forums/showthread.php?11474-Hex-Values-and-Trashbytes-in-B-W#post93598
+
+def make_pokemon_struct(generation):
+    """Make a pokemon struct class for the given generation
+    """
+    leaves_or_nature = {
+        4: BitStruct('shining_leaves',
+                Padding(2),
+                Flag('crown'),
+                Flag('leaf5'),
+                Flag('leaf4'),
+                Flag('leaf3'),
+                Flag('leaf2'),
+                Flag('leaf1'),
+            ),
+        5: ULInt8('nature_id'),
+    }[generation]
+
+    padding_or_hidden_ability = {
+        4: Padding(1),
+        5: Flag('hidden_ability'),
+    }[generation]
+
+    return Struct('pokemon_struct',
         # Header
         ULInt32('personality'),  # XXX aughgh http://bulbapedia.bulbagarden.net/wiki/Personality
         Padding(2),
@@ -604,8 +626,8 @@ pokemon_struct = Struct('pokemon_struct',
             Flag('triangle'),
             Flag('circle'),
         ),
-    Enum(
+        Enum(ULInt8('original_country'),
-        ULInt8('original_country'),
+            _unset = 0,
             jp=1,
             us=2,
             fr=3,
@@ -729,16 +751,9 @@ pokemon_struct = Struct('pokemon_struct',
             ),
             Flag('fateful_encounter'),
         ),
-    BitStruct('shining_leaves',
+        leaves_or_nature,
-        Padding(2),
+        padding_or_hidden_ability,
-        Flag('crown'),
+        Padding(1),
-        Flag('leaf5'),
-        Flag('leaf4'),
-        Flag('leaf3'),
-        Flag('leaf2'),
-        Flag('leaf1'),
-    ),
-    Padding(2),
         ULInt16('pt_egg_location_id'),
         ULInt16('pt_met_location_id'),
 
@@ -746,6 +761,7 @@ pokemon_struct = Struct('pokemon_struct',
         PokemonStringAdapter(String('nickname', 22)),
         Padding(1),
         Enum(ULInt8('original_version'),
+            _unset = 0,
             sapphire = 1,
             ruby = 2,
             emerald = 3,
@@ -789,7 +805,7 @@ pokemon_struct = Struct('pokemon_struct',
         DateAdapter(String('date_met', 3)),
         ULInt16('dp_egg_location_id'),
         ULInt16('dp_met_location_id'),
-    ULInt8('pokerus'),
+        ULInt8('pokerus'),  # Warning : Values changed in gen 5
         ULInt8('dppt_pokeball'),
         EmbeddedBitStruct(
             Enum(Flag('original_trainer_gender'),
@@ -812,4 +828,4 @@ pokemon_struct = Struct('pokemon_struct',
         ),
         ULInt8('hgss_pokeball'),
         Padding(1),
-)
+    )