X-Git-Url: http://git.veekun.com/zzz-pokedex.git/blobdiff_plain/45c94cc0640ae7af1308d7546a09cbce8e37f589..c87f91105dbce6befb0d4b68c1d6b9bb2a69e805:/pokedex/struct/__init__.py

diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py
index d1142d0..5b52539 100644
--- a/pokedex/struct/__init__.py
+++ b/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,6 +80,189 @@ 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."""
+        # See http://bulbapedia.bulbagarden.net/wiki/Personality#Shininess
+        # But don't see it too much, because the above is super over
+        # complicated.  Do this instead!
+        personality_msdw = self.structure.personality >> 16
+        personality_lsdw = self.structure.personality & 0xffff
+        return (
+            self.structure.original_trainer_id
+            ^ self.structure.original_trainer_secret_id
+            ^ personality_msdw
+            ^ 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._pokemon_form = session.query(tables.PokemonForm) \
+            .with_parent(self._pokemon) \
+            .filter_by(name=st.alternate_form) \
+            .one()
+        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.ItemGameIndex) \
+                .filter_by(generation_id = 4, game_index = 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.ItemGameIndex) \
+            .filter_by(generation_id = 4, game_index = 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.LocationGameIndex) \
+                .filter_by(generation_id = 4, game_index = egg_loc_id).one().location
+
+        self._met_location = session.query(tables.LocationGameIndex) \
+            .filter_by(generation_id = 4, game_index = met_loc_id).one().location
+
+    @property
+    def species(self):
+        # XXX forme!
+        return self._pokemon
+
+    @property
+    def species_form(self):
+        return self._pokemon_form
+
+    @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