diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index d119e82..07182bf 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -22,6 +22,28 @@ def pokemon_prng(seed): seed &= 0xFFFFFFFF yield seed >> 16 + +def _struct_proxy(name): + def getter(self): + return getattr(self.structure, name) + + def setter(self, value): + setattr(self.structure, name, value) + + return property(getter, setter) + +def _struct_frozenset_proxy(name): + def getter(self): + bitstruct = getattr(self.structure, name) + return frozenset(k for k, v in bitstruct.items() if v) + + def setter(self, value): + bitstruct = dict((k, True) for k in value) + setattr(self.structure, name, bitstruct) + + return property(getter, setter) + + class SaveFilePokemon(object): u"""Base class for an individual Pokémon, from the game's point of view. @@ -70,6 +92,8 @@ class SaveFilePokemon(object): if session: self.use_database_session(session) + else: + self._session = None @property def as_struct(self): @@ -91,6 +115,107 @@ class SaveFilePokemon(object): # Stuff back into a string, and done return struct.pack(struct_def, *shuffled) + def export(self): + """Exports the pokemon as a YAML/JSON-compatible dict + """ + st = self.structure + + result = dict( + species=dict(id=self.species.id, name=self.species.name), + ability=dict(id=self.ability.id, name=self.ability.name), + ) + + result['original trainer'] = dict( + id=self.original_trainer_id, + secret=self.original_trainer_secret_id, + name=unicode(self.original_trainer_name), + gender=self.original_trainer_gender + ) + + if self.form != self.species.default_form: + result['form'] = dict(id=self.form.id, name=self.form.form_name) + if self.held_item: + result['item'] = dict(id=self.item.id, name=self.item.name) + if self.exp: + result['exp'] = self.exp + if self.happiness: + result['happiness'] = self.happiness + if self.markings: + result['markings'] = sorted(self.markings) + if self.original_country and self.original_country != '_unset': + result['original country'] = self.original_country + if self.original_version and self.original_version != '_unset': + result['original version'] = self.original_version + if self.encounter_type and self.encounter_type != '_unset': + result['encounter type'] = self.encounter_type + if self.nickname: + result['nickname'] = unicode(self.nickname) + if self.egg_location: + result['egg location'] = dict(id=self.egg_location.id, + name=self.egg_location.name) + if self.met_location: + result['met location'] = dict(id=self.egg_location.id, + name=self.met_location.name) + if self.date_egg_received: + result['egg received'] = self.date_egg_received.isoformat() + if self.date_met: + result['date met'] = self.date_met.isoformat() + if self.pokerus: + result['pokerus data'] = self.pokerus + if self.pokeball: + result['pokeball'] = dict(id=self.pokeball.id, + name=self.pokeball.name) + if self.met_at_level: + result['met at level'] = self.met_at_level + + if not self.is_nicknamed: + result['not nicknamed'] = True + if self.is_egg: + result['is egg'] = True + if self.fateful_encounter: + result['fateful encounter'] = True + if self.gender != 'genderless': + result['gender'] = self.gender + + moves = result['moves'] = [] + for i, move_object in enumerate(self.moves, 1): + move = {} + if move_object: + move['id'] = move_object.id + move['name'] = move_object.name + pp = st['move%s_pp' % i] + if pp: + move['pp'] = pp + pp_up = st['move%s_pp_ups' % i] + if pp_up: + move['pp_up'] = pp_up + if move: + moves.append(move) + + effort = {} + genes = {} + contest_stats = {} + for pokemon_stat in self._pokemon.stats: + stat_identifier = pokemon_stat.stat.identifier.replace('-', '_') + if st['iv_' + stat_identifier]: + genes[stat_identifier] = st['iv_' + stat_identifier] + if st['effort_' + stat_identifier]: + effort[stat_identifier] = st['effort_' + stat_identifier] + for contest_stat in 'cool', 'beauty', 'cute', 'smart', 'tough', 'sheen': + if st['contest_' + contest_stat]: + contest_stats[contest_stat] = st['contest_' + contest_stat] + if effort: + result['effort'] = effort + if genes: + result['genes'] = genes + if contest_stats: + result['contest_stats'] = contest_stats + + ribbons = list(self.ribbons) + if ribbons: + result['ribbons'] = ribbons + return result + ### Delicious data @property def is_shiny(self): @@ -109,7 +234,7 @@ class SaveFilePokemon(object): def use_database_session(self, session): """Remembers the given database session, and prefetches a bunch of - database stuff. Gotta call this (or give it to `__init__`) before + database stuff. Gotta call this (or give session to `__init__`) before you use the database properties like `species`, etc. """ self._session = session @@ -150,11 +275,11 @@ class SaveFilePokemon(object): if self._pokemon: 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] + stat_identifier = pokemon_stat.stat.identifier.replace('-', '_') + gene = st['iv_' + stat_identifier] + exp = st['effort_' + stat_identifier] - if pokemon_stat.stat.name == u'HP': + if pokemon_stat.stat.identifier == u'hp': calc = calculated_hp else: calc = calculated_stat @@ -285,6 +410,58 @@ class SaveFilePokemon(object): self.structure.move4_pp, ) + @property + def markings(self): + return frozenset(k for k, v in self.structure.markings.items() if v) + + @markings.setter + def markings(self, value): + self.structure.markings = dict((k, True) for k in value) + + original_trainer_id = _struct_proxy('original_trainer_id') + original_trainer_secret_id = _struct_proxy('original_trainer_secret_id') + original_trainer_name = _struct_proxy('original_trainer_name') + exp = _struct_proxy('exp') + happiness = _struct_proxy('happiness') + original_country = _struct_proxy('original_country') + is_nicknamed = _struct_proxy('is_nicknamed') + is_egg = _struct_proxy('is_egg') + fateful_encounter = _struct_proxy('fateful_encounter') + gender = _struct_proxy('gender') + original_version = _struct_proxy('original_version') + date_egg_received = _struct_proxy('date_egg_received') + date_met = _struct_proxy('date_met') + pokerus = _struct_proxy('pokerus') + met_at_level = _struct_proxy('met_at_level') + original_trainer_gender = _struct_proxy('original_trainer_gender') + encounter_type = _struct_proxy('encounter_type') + + markings = _struct_frozenset_proxy('markings') + sinnoh_ribbons = _struct_frozenset_proxy('sinnoh_ribbons') + hoenn_ribbons = _struct_frozenset_proxy('hoenn_ribbons') + sinnoh_contest_ribbons = _struct_frozenset_proxy('sinnoh_contest_ribbons') + + @property + def ribbons(self): + return frozenset( + self.sinnoh_ribbons | + self.hoenn_ribbons | + self.sinnoh_contest_ribbons) + # XXX: ribbons setter + + @property + def nickname(self): + return self.structure.nickname + + @nickname.setter + def nickname(self, value): + self.structure.nickname = value + self.structure.is_nicknamed = True + + @nickname.deleter + def nickname(self, value): + self.structure.nickname = '' + self.structure.is_nicknamed = False ### Utility methods @@ -342,11 +519,31 @@ class SaveFilePokemon(object): return + def _reset(self): + """Update self with modified pokemon_struct + + Rebuilds the blob; recalculates checksum; if a session was set, + re-fetches the DB objects. + """ + self.blob = self.pokemon_struct.build(self.structure) + self.structure = self.pokemon_struct.parse(self.blob) + checksum = sum(struct.unpack('H' * 0x40, self.blob[8:0x88])) & 0xffff + self.structure.checksum = checksum + self.blob = self.blob[:6] + struct.pack('H', checksum) + self.blob[8:] + if self._session: + self.use_database_session(self._session) + class SaveFilePokemonGen4(SaveFilePokemon): generation_id = 4 pokemon_struct = make_pokemon_struct(generation=generation_id) + def export(self): + result = super(SaveFilePokemonGen5, self).export() + if any(self.shiny_leaves): + result['shiny leaves'] = self.shiny_leaves + return result + @property def shiny_leaves(self): return ( @@ -355,8 +552,21 @@ class SaveFilePokemonGen4(SaveFilePokemon): self.structure.shining_leaves.leaf3, self.structure.shining_leaves.leaf4, self.structure.shining_leaves.leaf5, + self.structure.shining_leaves.crown, ) + @shiny_leaves.setter + def shiny_leaves(self, new_values): + ( + 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, + self.structure.shining_leaves.crown, + ) = new_values + self._reset() + class SaveFilePokemonGen5(SaveFilePokemon): generation_id = 5 @@ -371,10 +581,22 @@ class SaveFilePokemonGen5(SaveFilePokemon): self._nature = session.query(tables.Nature) \ .filter_by(game_index = st.nature_id).one() + def export(self): + result = super(SaveFilePokemonGen5, self).export() + result['nature'] = dict(id=self.nature.id, name=self.nature.name) + return result + + # XXX: Ability setter must set hidden ability flag + @property def nature(self): return self._nature + @nature.setter + def nature(self, new_nature): + self.structure.nature_id = int(new_nature.game_index) + self._reset() + save_file_pokemon_classes = { 4: SaveFilePokemonGen4, diff --git a/pokedex/struct/_pokemon_struct.py b/pokedex/struct/_pokemon_struct.py index 755683c..f9590d3 100644 --- a/pokedex/struct/_pokemon_struct.py +++ b/pokedex/struct/_pokemon_struct.py @@ -792,7 +792,7 @@ def make_pokemon_struct(generation): ULInt8('move3_pp_ups'), ULInt8('move4_pp_ups'), - LittleEndianBitStruct('ivs', + Embed(LittleEndianBitStruct('ivs', Flag('is_nicknamed'), Flag('is_egg'), BitField('iv_special_defense', 5), @@ -801,7 +801,7 @@ def make_pokemon_struct(generation): BitField('iv_defense', 5), BitField('iv_attack', 5), BitField('iv_hp', 5), - ), + )), LittleEndianBitStruct('hoenn_ribbons', Flag('world_ribbon'), Flag('earth_ribbon'), @@ -836,7 +836,7 @@ def make_pokemon_struct(generation): Flag('cool_ribbon_super'), Flag('cool_ribbon'), ), - EmbeddedBitStruct( + Embed(EmbeddedBitStruct( PokemonFormAdapter(BitField('alternate_form', 5)), Enum(BitField('gender', 2), genderless = 2, @@ -844,7 +844,7 @@ def make_pokemon_struct(generation): female = 1, ), Flag('fateful_encounter'), - ), + )), leaves_or_nature, padding_or_hidden_ability, Padding(1),