Adapt the Pokémon savefile struct to the database.

pokedex/struct/__init__.py

@@ -10,7 +10,9 @@ derived.
 
 import struct
 
-from pokedex.util import permutations
+from pokedex.db import tables
+from pokedex.formulae import calculated_hp, calculated_stat
+from pokedex.util import namedtuple, permutations
 from pokedex.struct._pokemon_struct import pokemon_struct
 
 def pokemon_prng(seed):
@@ -28,12 +30,16 @@ class SaveFilePokemon(object):
     (also used by pokesav), and something vaguely intelligible.
     """
 
+    Stat = namedtuple('Stat', ['stat', 'base', 'gene', 'exp', 'calc'])
+
     def __init__(self, blob, encrypted=False):
         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.
         """
 
         if encrypted:
@@ -54,7 +60,6 @@ class SaveFilePokemon(object):
 
         self.structure = pokemon_struct.parse(self.blob)
 
-
     @property
     def as_struct(self):
         u"""Returns a decrypted struct, aka .pkm file."""
@@ -75,9 +80,7 @@ class SaveFilePokemon(object):
         # Stuff back into a string, and done
         return struct.pack(struct_def, *shuffled)
 
-
     ### Delicious data
-
     @property
     def is_shiny(self):
         u"""Returns true iff this Pokémon is shiny."""
@@ -93,6 +96,165 @@ class SaveFilePokemon(object):
             ^ personality_lsdw
         ) < 8
 
+    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
+        like `species`, etc.
+        """
+        self._session = session
+
+        st = self.structure
+        self._pokemon = session.query(tables.Pokemon).get(st.national_id)
+        self._ability = self._session.query(tables.Ability).get(st.ability_id)
+
+        growth_rate = self._pokemon.evolution_chain.growth_rate
+        self._experience_rung = session.query(tables.Experience) \
+            .filter(tables.Experience.growth_rate == growth_rate) \
+            .filter(tables.Experience.experience <= st.exp) \
+            .order_by(tables.Experience.level.desc()) \
+            [0]
+        level = self._experience_rung.level
+
+        self._next_experience_rung = None
+        if level < 100:
+            self._next_experience_rung = session.query(tables.Experience) \
+                .filter(tables.Experience.growth_rate == growth_rate) \
+                .filter(tables.Experience.level == level + 1) \
+                .one()
+
+        self._held_item = None
+        if st.held_item_id:
+            self._held_item = session.query(tables.ItemInternalID) \
+                .filter_by(generation_id = 4, internal_id = st.held_item_id).one().item
+
+        self._stats = []
+        for pokemon_stat in self._pokemon.stats:
+            structure_name = pokemon_stat.stat.name.lower().replace(' ', '_')
+            gene = st.ivs['iv_' + structure_name]
+            exp  = st['effort_' + structure_name]
+
+            if pokemon_stat.stat.name == u'HP':
+                calc = calculated_hp
+            else:
+                calc = calculated_stat
+
+            stat_tup = self.Stat(
+                stat = pokemon_stat.stat,
+                base = pokemon_stat.base_stat,
+                gene = gene,
+                exp  = exp,
+                calc = calc(
+                    pokemon_stat.base_stat,
+                    level = level,
+                    iv = gene,
+                    effort = exp,
+                ),
+            )
+
+            self._stats.append(stat_tup)
+
+
+        move_ids = (
+            self.structure.move1_id,
+            self.structure.move2_id,
+            self.structure.move3_id,
+            self.structure.move4_id,
+        )
+        move_rows = self._session.query(tables.Move).filter(tables.Move.id.in_(move_ids))
+        moves_dict = dict((move.id, move) for move in move_rows)
+
+        self._moves = [moves_dict.get(move_id, None) for move_id in move_ids]
+
+        if st.hgss_pokeball >= 17:
+            pokeball_id = st.hgss_pokeball - 17 + 492
+        else:
+            pokeball_id = st.dppt_pokeball
+        self._pokeball = session.query(tables.ItemInternalID) \
+            .filter_by(generation_id = 4, internal_id = pokeball_id).one().item
+
+        egg_loc_id = st.pt_egg_location_id or st.dp_egg_location_id
+        met_loc_id = st.pt_met_location_id or st.dp_met_location_id
+
+        self._egg_location = None
+        if egg_loc_id:
+            self._egg_location = session.query(tables.LocationInternalID) \
+                .filter_by(generation_id = 4, internal_id = egg_loc_id).one().location
+
+        self._met_location = session.query(tables.LocationInternalID) \
+            .filter_by(generation_id = 4, internal_id = met_loc_id).one().location
+
+    @property
+    def species(self):
+        # XXX forme!
+        return self._pokemon
+
+    @property
+    def pokeball(self):
+        return self._pokeball
+
+    @property
+    def egg_location(self):
+        return self._egg_location
+
+    @property
+    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
+
+    @property
+    def exp_to_next(self):
+        if self._next_experience_rung:
+            return self._next_experience_rung.experience - self.structure.exp
+        else:
+            return 0
+
+    @property
+    def progress_to_next(self):
+        if self._next_experience_rung:
+            return 1.0 \
+                * (self.structure.exp - self._experience_rung.experience) \
+                / (self._next_experience_rung.experience - self._experience_rung.experience)
+        else:
+            return 0.0
+
+    @property
+    def ability(self):
+        return self._ability
+
+    @property
+    def held_item(self):
+        return self._held_item
+
+    @property
+    def stats(self):
+        return self._stats
+
+    @property
+    def moves(self):
+        return self._moves
+
+    @property
+    def move_pp(self):
+        return (
+            self.structure.move1_pp,
+            self.structure.move2_pp,
+            self.structure.move3_pp,
+            self.structure.move4_pp,
+        )
+
 
     ### Utility methods
 

pokedex/struct/_pokemon_struct.py

@@ -527,6 +527,58 @@ class DateAdapter(Adapter):
         y, m, d = obj.year - 2000, obj.month, obj.day
         return ''.join(chr(n) for n in (y, m, d))
 
+class PokemonFormAdapter(Adapter):
+    """Converts form ids to form names, and vice versa."""
+    pokemon_forms = {
+        # Unown
+        201: 'abcdefghijklmnopqrstuvwxyz!?',
+
+        # Deoxys
+        386: ['normal', 'attack', 'defense', 'speed'],
+
+        # Burmy and Wormadam
+        412: ['plant', 'sandy', 'trash'],
+        413: ['plant', 'sandy', 'trash'],
+
+        # Shellos and Gastrodon
+        422: ['west', 'east'],
+        423: ['west', 'east'],
+
+        # Rotom
+        479: ['normal', 'heat', 'wash', 'frost', 'fan', 'cut'],
+
+        # Giratina
+        487: ['altered', 'origin'],
+
+        # Shaymin
+        492: ['land', 'sky'],
+
+        # Arceus
+        493: [
+            'normal', 'fighting', 'flying', 'poison', 'ground', 'rock',
+            'bug', 'ghost', 'steel', 'fire', 'water', 'grass',
+            'thunder', 'psychic', 'ice', 'dragon', 'dark', '???',
+        ],
+    }
+
+    def _decode(self, obj, context):
+        try:
+            forms = self.pokemon_forms[ context['national_id'] ]
+        except KeyError:
+            return None
+
+        return forms[obj >> 3]
+
+    def _encode(self, obj, context):
+        try:
+            forms = self.pokemon_forms[ context['national_id'] ]
+        except KeyError:
+            return None
+
+        return forms.index(obj) << 3
+
+
+
 # And here we go.
 # Docs: http://projectpokemon.org/wiki/Pokemon_NDS_Structure
 pokemon_struct = Struct('pokemon_struct',
@@ -669,7 +721,7 @@ pokemon_struct = Struct('pokemon_struct',
         Flag('cool_ribbon'),
     ),
     EmbeddedBitStruct(
-        BitField('alternate_form', 5),
+        PokemonFormAdapter(BitField('alternate_form', 5)),
         Enum(BitField('gender', 2),
             genderless = 2,
             male = 0,
@@ -747,14 +799,16 @@ pokemon_struct = Struct('pokemon_struct',
         BitField('met_at_level', 7),
     ),
     Enum(ULInt8('encounter_type'),
-        special = 0,
+        special = 0,        # egg; pal park; event; honey tree; shaymin
-        grass = 2,
+        grass = 2,          # or darkrai
         dialga_palkia = 4,
-        cave = 5,  # or hall of origin
+        cave = 5,           # or giratina or hall of origin
         water = 7,
         building = 9,
-        safari_zone = 10,
+        safari_zone = 10,   # includes great marsh
-        gift = 12,
+        gift = 12,          # starter; fossil; ingame trade?
+        # distortion_world = ???,
+        hgss_gift = 24,     # starter; fossil; bebe's eevee  (pt only??)
     ),
     ULInt8('hgss_pokeball'),
     Padding(1),