From 56d5c3e9536131f24b60b764753bb18d3c3b7a02 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sat, 9 Jun 2012 21:08:36 +0200 Subject: [PATCH 01/23] Add game_index to Nature --- pokedex/data/csv/natures.csv | 52 ++++++++++++++++++------------------ pokedex/db/tables.py | 2 ++ 2 files changed, 28 insertions(+), 26 deletions(-) diff --git a/pokedex/data/csv/natures.csv b/pokedex/data/csv/natures.csv index d5864af6d..81a27fe47 100644 --- a/pokedex/data/csv/natures.csv +++ b/pokedex/data/csv/natures.csv @@ -1,26 +1,26 @@ -id,identifier,decreased_stat_id,increased_stat_id,hates_flavor_id,likes_flavor_id -1,hardy,2,2,1,1 -2,bold,2,3,1,5 -3,modest,2,4,1,2 -4,calm,2,5,1,4 -5,timid,2,6,1,3 -6,lonely,3,2,5,1 -7,docile,3,3,5,5 -8,mild,3,4,5,2 -9,gentle,3,5,5,4 -10,hasty,3,6,5,3 -11,adamant,4,2,2,1 -12,impish,4,3,2,5 -13,bashful,4,4,2,2 -14,careful,4,5,2,4 -15,rash,5,4,4,2 -16,jolly,4,6,2,3 -17,naughty,5,2,4,1 -18,lax,5,3,4,5 -19,quirky,5,5,4,4 -20,naive,5,6,4,3 -21,brave,6,2,3,1 -22,relaxed,6,3,3,5 -23,quiet,6,4,3,2 -24,sassy,6,5,3,4 -25,serious,6,6,3,3 +id,identifier,decreased_stat_id,increased_stat_id,hates_flavor_id,likes_flavor_id,game_index +1,hardy,2,2,1,1,0 +2,bold,2,3,1,5,5 +3,modest,2,4,1,2,15 +4,calm,2,5,1,4,20 +5,timid,2,6,1,3,10 +6,lonely,3,2,5,1,1 +7,docile,3,3,5,5,6 +8,mild,3,4,5,2,16 +9,gentle,3,5,5,4,21 +10,hasty,3,6,5,3,11 +11,adamant,4,2,2,1,3 +12,impish,4,3,2,5,8 +13,bashful,4,4,2,2,18 +14,careful,4,5,2,4,23 +15,rash,5,4,4,2,19 +16,jolly,4,6,2,3,13 +17,naughty,5,2,4,1,4 +18,lax,5,3,4,5,9 +19,quirky,5,5,4,4,24 +20,naive,5,6,4,3,14 +21,brave,6,2,3,1,2 +22,relaxed,6,3,3,5,7 +23,quiet,6,4,3,2,17 +24,sassy,6,5,3,4,22 +25,serious,6,6,3,3,12 diff --git a/pokedex/db/tables.py b/pokedex/db/tables.py index ad02a4b0f..043bfacde 100644 --- a/pokedex/db/tables.py +++ b/pokedex/db/tables.py @@ -1392,6 +1392,8 @@ class Nature(TableBase): info=dict(description=u"ID of the Berry flavor the Pokémon hates (if likes_flavor_id is the same, the effects cancel out)")) likes_flavor_id = Column(Integer, ForeignKey('contest_types.id'), nullable=False, info=dict(description=u"ID of the Berry flavor the Pokémon likes (if hates_flavor_id is the same, the effects cancel out)")) + game_index = Column(Integer, unique=True, nullable=False, + info=dict(description=u"Internal game ID of the nature")) @property def is_neutral(self): From 198a1fac8d5e04092a7791b6a05276db26f16dee Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 10 Jun 2012 23:01:15 +0200 Subject: [PATCH 02/23] 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 | 225 +++++++++----- pokedex/struct/_pokemon_struct.py | 476 +++++++++++++++--------------- 2 files changed, 389 insertions(+), 312 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index b6a3e736a..d119e8231 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -13,7 +13,7 @@ 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,43 +22,54 @@ 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` """ - if encrypted: - # Decrypt it. - # Interpret as one word (pid), followed by a bunch of shorts - struct_def = "I" + "H" * ((len(blob) - 4) / 2) - shuffled = list( struct.unpack(struct_def, blob) ) + 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 + struct_def = "I" + "H" * ((len(blob) - 4) / 2) + shuffled = list( struct.unpack(struct_def, blob) ) + + # Apply standard Pokémon decryption, undo the block shuffling, and + # done + self.reciprocal_crypt(shuffled) + words = self.shuffle_chunks(shuffled, reverse=True) + self.blob = struct.pack(struct_def, *words) - # Apply standard Pokémon decryption, undo the block shuffling, and - # done - self.reciprocal_crypt(shuffled) - words = self.shuffle_chunks(shuffled, reverse=True) - self.blob = struct.pack(struct_def, *words) + else: + # Already decrypted + self.blob = blob + self.structure = self.pokemon_struct.parse(self.blob) else: - # Already decrypted - self.blob = blob + self.structure = self.pokemon_struct.parse('\0' * (32 * 4 + 8)) - self.structure = pokemon_struct.parse(self.blob) + if session: + self.use_database_session(session) @property def as_struct(self): @@ -98,64 +109,72 @@ def is_shiny(self): 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. + database stuff. Gotta call this (or give it to `__init__`) 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() + + 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(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 - 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) \ + 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.level == level + 1) \ - .one() + .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) + 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] + + 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) + else: + self._stats = [0] * 6 move_ids = ( @@ -171,10 +190,13 @@ def use_database_session(self, session): if st.hgss_pokeball >= 17: pokeball_id = st.hgss_pokeball - 17 + 492 - else: + elif st.dppt_pokeball: pokeball_id = st.dppt_pokeball - self._pokeball = session.query(tables.ItemGameIndex) \ - .filter_by(generation_id = 4, game_index = pokeball_id).one().item + else: + pokeball_id = None + if pokeball_id: + 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 @@ -182,15 +204,25 @@ def use_database_session(self, session): 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 - self._met_location = session.query(tables.LocationGameIndex) \ - .filter_by(generation_id = 4, game_index = met_loc_id).one().location + if met_loc_id: + self._met_location = session.query(tables.LocationGameIndex) \ + .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 + return self._pokemon_form.species + + @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 @@ def egg_location(self): 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 @@ def reciprocal_crypt(cls, words): 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, +} diff --git a/pokedex/struct/_pokemon_struct.py b/pokedex/struct/_pokemon_struct.py index 79ff0e0f4..cf71c2437 100644 --- a/pokedex/struct/_pokemon_struct.py +++ b/pokedex/struct/_pokemon_struct.py @@ -578,238 +578,254 @@ def _encode(self, obj, context): return forms.index(obj) << 3 - -# And here we go. # Docs: http://projectpokemon.org/wiki/Pokemon_NDS_Structure -pokemon_struct = Struct('pokemon_struct', - # Header - ULInt32('personality'), # XXX aughgh http://bulbapedia.bulbagarden.net/wiki/Personality - Padding(2), - ULInt16('checksum'), # XXX should be checked or calculated - - # Block A - ULInt16('national_id'), - ULInt16('held_item_id'), - ULInt16('original_trainer_id'), - ULInt16('original_trainer_secret_id'), - ULInt32('exp'), - ULInt8('happiness'), - ULInt8('ability_id'), # XXX needs to match personality + species - BitStruct('markings', +# 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), - Flag('diamond'), - Flag('star'), - Flag('heart'), - Flag('square'), - Flag('triangle'), - Flag('circle'), - ), - Enum( - ULInt8('original_country'), - jp=1, - us=2, - fr=3, - it=4, - de=5, - es=7, - kr=8, - ), - - # XXX sum cannot surpass 510 - ULInt8('effort_hp'), - ULInt8('effort_attack'), - ULInt8('effort_defense'), - ULInt8('effort_speed'), - ULInt8('effort_special_attack'), - ULInt8('effort_special_defense'), - - ULInt8('contest_cool'), - ULInt8('contest_beauty'), - ULInt8('contest_cute'), - ULInt8('contest_smart'), - ULInt8('contest_tough'), - ULInt8('contest_sheen'), - - LittleEndianBitStruct('sinnoh_ribbons', + ULInt16('checksum'), # XXX should be checked or calculated + + # Block A + ULInt16('national_id'), + ULInt16('held_item_id'), + ULInt16('original_trainer_id'), + ULInt16('original_trainer_secret_id'), + ULInt32('exp'), + ULInt8('happiness'), + ULInt8('ability_id'), # XXX needs to match personality + species + BitStruct('markings', + Padding(2), + Flag('diamond'), + Flag('star'), + Flag('heart'), + Flag('square'), + Flag('triangle'), + Flag('circle'), + ), + Enum(ULInt8('original_country'), + _unset = 0, + jp=1, + us=2, + fr=3, + it=4, + de=5, + es=7, + kr=8, + ), + + # XXX sum cannot surpass 510 + ULInt8('effort_hp'), + ULInt8('effort_attack'), + ULInt8('effort_defense'), + ULInt8('effort_speed'), + ULInt8('effort_special_attack'), + ULInt8('effort_special_defense'), + + ULInt8('contest_cool'), + ULInt8('contest_beauty'), + ULInt8('contest_cute'), + ULInt8('contest_smart'), + ULInt8('contest_tough'), + ULInt8('contest_sheen'), + + LittleEndianBitStruct('sinnoh_ribbons', + Padding(4), + Flag('premier_ribbon'), + Flag('classic_ribbon'), + Flag('carnival_ribbon'), + Flag('festival_ribbon'), + Flag('blue_ribbon'), + Flag('green_ribbon'), + Flag('red_ribbon'), + Flag('legend_ribbon'), + Flag('history_ribbon'), + Flag('record_ribbon'), + Flag('footprint_ribbon'), + Flag('gorgeous_royal_ribbon'), + Flag('royal_ribbon'), + Flag('gorgeous_ribbon'), + Flag('smile_ribbon'), + Flag('snooze_ribbon'), + Flag('relax_ribbon'), + Flag('careless_ribbon'), + Flag('downcast_ribbon'), + Flag('shock_ribbon'), + Flag('alert_ribbon'), + Flag('world_ability_ribbon'), + Flag('pair_ability_ribbon'), + Flag('multi_ability_ribbon'), + Flag('double_ability_ribbon'), + Flag('great_ability_ribbon'), + Flag('ability_ribbon'), + Flag('sinnoh_champ_ribbon'), + ), + + # Block B + ULInt16('move1_id'), + ULInt16('move2_id'), + ULInt16('move3_id'), + ULInt16('move4_id'), + ULInt8('move1_pp'), + ULInt8('move2_pp'), + ULInt8('move3_pp'), + ULInt8('move4_pp'), + ULInt8('move1_pp_ups'), + ULInt8('move2_pp_ups'), + ULInt8('move3_pp_ups'), + ULInt8('move4_pp_ups'), + + LittleEndianBitStruct('ivs', + Flag('is_nicknamed'), + Flag('is_egg'), + BitField('iv_special_defense', 5), + BitField('iv_special_attack', 5), + BitField('iv_speed', 5), + BitField('iv_defense', 5), + BitField('iv_attack', 5), + BitField('iv_hp', 5), + ), + LittleEndianBitStruct('hoenn_ribbons', + Flag('world_ribbon'), + Flag('earth_ribbon'), + Flag('national_ribbon'), + Flag('country_ribbon'), + Flag('sky_ribbon'), + Flag('land_ribbon'), + Flag('marine_ribbon'), + Flag('effort_ribbon'), + Flag('artist_ribbon'), + Flag('victory_ribbon'), + Flag('winning_ribbon'), + Flag('champion_ribbon'), + Flag('tough_ribbon_master'), + Flag('tough_ribbon_hyper'), + Flag('tough_ribbon_super'), + Flag('tough_ribbon'), + Flag('smart_ribbon_master'), + Flag('smart_ribbon_hyper'), + Flag('smart_ribbon_super'), + Flag('smart_ribbon'), + Flag('cute_ribbon_master'), + Flag('cute_ribbon_hyper'), + Flag('cute_ribbon_super'), + Flag('cute_ribbon'), + Flag('beauty_ribbon_master'), + Flag('beauty_ribbon_hyper'), + Flag('beauty_ribbon_super'), + Flag('beauty_ribbon'), + Flag('cool_ribbon_master'), + Flag('cool_ribbon_hyper'), + Flag('cool_ribbon_super'), + Flag('cool_ribbon'), + ), + EmbeddedBitStruct( + PokemonFormAdapter(BitField('alternate_form', 5)), + Enum(BitField('gender', 2), + genderless = 2, + male = 0, + female = 1, + ), + Flag('fateful_encounter'), + ), + leaves_or_nature, + padding_or_hidden_ability, + Padding(1), + ULInt16('pt_egg_location_id'), + ULInt16('pt_met_location_id'), + + # Block C + PokemonStringAdapter(String('nickname', 22)), + Padding(1), + Enum(ULInt8('original_version'), + _unset = 0, + sapphire = 1, + ruby = 2, + emerald = 3, + firered = 4, + leafgreen = 5, + heartgold = 7, + soulsilver = 8, + diamond = 10, + pearl = 11, + platinum = 12, + orre = 15, + ), + LittleEndianBitStruct('sinnoh_contest_ribbons', + Padding(12), + Flag('tough_ribbon_master'), + Flag('tough_ribbon_ultra'), + Flag('tough_ribbon_great'), + Flag('tough_ribbon'), + Flag('smart_ribbon_master'), + Flag('smart_ribbon_ultra'), + Flag('smart_ribbon_great'), + Flag('smart_ribbon'), + Flag('cute_ribbon_master'), + Flag('cute_ribbon_ultra'), + Flag('cute_ribbon_great'), + Flag('cute_ribbon'), + Flag('beauty_ribbon_master'), + Flag('beauty_ribbon_ultra'), + Flag('beauty_ribbon_great'), + Flag('beauty_ribbon'), + Flag('cool_ribbon_master'), + Flag('cool_ribbon_ultra'), + Flag('cool_ribbon_great'), + Flag('cool_ribbon'), + ), Padding(4), - Flag('premier_ribbon'), - Flag('classic_ribbon'), - Flag('carnival_ribbon'), - Flag('festival_ribbon'), - Flag('blue_ribbon'), - Flag('green_ribbon'), - Flag('red_ribbon'), - Flag('legend_ribbon'), - Flag('history_ribbon'), - Flag('record_ribbon'), - Flag('footprint_ribbon'), - Flag('gorgeous_royal_ribbon'), - Flag('royal_ribbon'), - Flag('gorgeous_ribbon'), - Flag('smile_ribbon'), - Flag('snooze_ribbon'), - Flag('relax_ribbon'), - Flag('careless_ribbon'), - Flag('downcast_ribbon'), - Flag('shock_ribbon'), - Flag('alert_ribbon'), - Flag('world_ability_ribbon'), - Flag('pair_ability_ribbon'), - Flag('multi_ability_ribbon'), - Flag('double_ability_ribbon'), - Flag('great_ability_ribbon'), - Flag('ability_ribbon'), - Flag('sinnoh_champ_ribbon'), - ), - - # Block B - ULInt16('move1_id'), - ULInt16('move2_id'), - ULInt16('move3_id'), - ULInt16('move4_id'), - ULInt8('move1_pp'), - ULInt8('move2_pp'), - ULInt8('move3_pp'), - ULInt8('move4_pp'), - ULInt8('move1_pp_ups'), - ULInt8('move2_pp_ups'), - ULInt8('move3_pp_ups'), - ULInt8('move4_pp_ups'), - - LittleEndianBitStruct('ivs', - Flag('is_nicknamed'), - Flag('is_egg'), - BitField('iv_special_defense', 5), - BitField('iv_special_attack', 5), - BitField('iv_speed', 5), - BitField('iv_defense', 5), - BitField('iv_attack', 5), - BitField('iv_hp', 5), - ), - LittleEndianBitStruct('hoenn_ribbons', - Flag('world_ribbon'), - Flag('earth_ribbon'), - Flag('national_ribbon'), - Flag('country_ribbon'), - Flag('sky_ribbon'), - Flag('land_ribbon'), - Flag('marine_ribbon'), - Flag('effort_ribbon'), - Flag('artist_ribbon'), - Flag('victory_ribbon'), - Flag('winning_ribbon'), - Flag('champion_ribbon'), - Flag('tough_ribbon_master'), - Flag('tough_ribbon_hyper'), - Flag('tough_ribbon_super'), - Flag('tough_ribbon'), - Flag('smart_ribbon_master'), - Flag('smart_ribbon_hyper'), - Flag('smart_ribbon_super'), - Flag('smart_ribbon'), - Flag('cute_ribbon_master'), - Flag('cute_ribbon_hyper'), - Flag('cute_ribbon_super'), - Flag('cute_ribbon'), - Flag('beauty_ribbon_master'), - Flag('beauty_ribbon_hyper'), - Flag('beauty_ribbon_super'), - Flag('beauty_ribbon'), - Flag('cool_ribbon_master'), - Flag('cool_ribbon_hyper'), - Flag('cool_ribbon_super'), - Flag('cool_ribbon'), - ), - EmbeddedBitStruct( - PokemonFormAdapter(BitField('alternate_form', 5)), - Enum(BitField('gender', 2), - genderless = 2, - male = 0, - female = 1, + + # Block D + PokemonStringAdapter(String('original_trainer_name', 16)), + DateAdapter(String('date_egg_received', 3)), + DateAdapter(String('date_met', 3)), + ULInt16('dp_egg_location_id'), + ULInt16('dp_met_location_id'), + ULInt8('pokerus'), # Warning : Values changed in gen 5 + ULInt8('dppt_pokeball'), + EmbeddedBitStruct( + Enum(Flag('original_trainer_gender'), + male = False, + female = True, + ), + BitField('met_at_level', 7), ), - Flag('fateful_encounter'), - ), - BitStruct('shining_leaves', - Padding(2), - Flag('crown'), - Flag('leaf5'), - Flag('leaf4'), - Flag('leaf3'), - Flag('leaf2'), - Flag('leaf1'), - ), - Padding(2), - ULInt16('pt_egg_location_id'), - ULInt16('pt_met_location_id'), - - # Block C - PokemonStringAdapter(String('nickname', 22)), - Padding(1), - Enum(ULInt8('original_version'), - sapphire = 1, - ruby = 2, - emerald = 3, - firered = 4, - leafgreen = 5, - heartgold = 7, - soulsilver = 8, - diamond = 10, - pearl = 11, - platinum = 12, - orre = 15, - ), - LittleEndianBitStruct('sinnoh_contest_ribbons', - Padding(12), - Flag('tough_ribbon_master'), - Flag('tough_ribbon_ultra'), - Flag('tough_ribbon_great'), - Flag('tough_ribbon'), - Flag('smart_ribbon_master'), - Flag('smart_ribbon_ultra'), - Flag('smart_ribbon_great'), - Flag('smart_ribbon'), - Flag('cute_ribbon_master'), - Flag('cute_ribbon_ultra'), - Flag('cute_ribbon_great'), - Flag('cute_ribbon'), - Flag('beauty_ribbon_master'), - Flag('beauty_ribbon_ultra'), - Flag('beauty_ribbon_great'), - Flag('beauty_ribbon'), - Flag('cool_ribbon_master'), - Flag('cool_ribbon_ultra'), - Flag('cool_ribbon_great'), - Flag('cool_ribbon'), - ), - Padding(4), - - # Block D - PokemonStringAdapter(String('original_trainer_name', 16)), - DateAdapter(String('date_egg_received', 3)), - DateAdapter(String('date_met', 3)), - ULInt16('dp_egg_location_id'), - ULInt16('dp_met_location_id'), - ULInt8('pokerus'), - ULInt8('dppt_pokeball'), - EmbeddedBitStruct( - Enum(Flag('original_trainer_gender'), - male = False, - female = True, + Enum(ULInt8('encounter_type'), + special = 0, # egg; pal park; event; honey tree; shaymin + grass = 2, # or darkrai + dialga_palkia = 4, + cave = 5, # or giratina or hall of origin + water = 7, + building = 9, + safari_zone = 10, # includes great marsh + gift = 12, # starter; fossil; ingame trade? + # distortion_world = ???, + hgss_gift = 24, # starter; fossil; bebe's eevee (pt only??) ), - BitField('met_at_level', 7), - ), - Enum(ULInt8('encounter_type'), - special = 0, # egg; pal park; event; honey tree; shaymin - grass = 2, # or darkrai - dialga_palkia = 4, - cave = 5, # or giratina or hall of origin - water = 7, - building = 9, - safari_zone = 10, # includes great marsh - gift = 12, # starter; fossil; ingame trade? - # distortion_world = ???, - hgss_gift = 24, # starter; fossil; bebe's eevee (pt only??) - ), - ULInt8('hgss_pokeball'), - Padding(1), -) + ULInt8('hgss_pokeball'), + Padding(1), + ) From bc672dde83f94723273d0ddbda9ecd17db64b027 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 17 Jun 2012 21:37:27 +0200 Subject: [PATCH 03/23] Add character map for generation 5 --- pokedex/struct/_pokemon_struct.py | 94 +++++++++++++++++++++++++++---- 1 file changed, 83 insertions(+), 11 deletions(-) diff --git a/pokedex/struct/_pokemon_struct.py b/pokedex/struct/_pokemon_struct.py index cf71c2437..40176e5ea 100644 --- a/pokedex/struct/_pokemon_struct.py +++ b/pokedex/struct/_pokemon_struct.py @@ -18,7 +18,7 @@ # - personality indirectly influences IVs due to PRNG use # The entire gen 4 character table: -character_table = { +character_table_gen4 = { 0x0002: u'ぁ', 0x0003: u'あ', 0x0004: u'ぃ', @@ -465,10 +465,64 @@ 0x25bd: u'\r', } -# And the reverse dict, used with str.translate() -inverse_character_table = dict() -for in_, out in character_table.iteritems(): - inverse_character_table[ord(out)] = in_ +# Generation 5 uses UCS-16, with a few exceptions +character_table_gen5 = { + # Here nintendo just didn't do their homework: + 0x247d: u'☂', + 0x247b: u'☁', + 0x247a: u'☀', + 0x2479: u'♪', + 0x2478: u'◇', + 0x2477: u'△', + 0x2476: u'□', + 0x2475: u'○', + 0x2474: u'◎', + 0x2473: u'★', + 0x2472: u'♦', + 0x2471: u'♥', + 0x2470: u'♣', + 0x246f: u'♠', + 0x246e: u'♀', + 0x246d: u'♂', + 0x246c: u'…', + 0x2468: u'÷', + 0x2467: u'×', + 0x21d4: u'⤴', + 0x2200: u'⤵', + + # These aren't direct equivalents, but better than nothing: + 0x0024: u'$', # pokémoney sign + 0x21d2: u'☹', # frowny face + 0x2203: u'ℤ', # ZZ ligature + 0x2227: u'☺', # smiling face + 0x2228: u'😁', # grinning face + 0xffe2: u'😭', # hurt face + + # The following duplicates & weird characters get to keep their positions + # ①..⑦ + # 0x2460: halfwidth smiling face + # 0x2461: grinning face + # 0x2462: hurt face + # 0x2463: frowny face + # 0x2464: ⤴ + # 0x2465: ⤵ + # 0x2466: ZZ ligature + # ⑩..⑫ + # 0x2469: superscript er + # 0x246a: superscript re + # 0x246b: superscript r + # ⑾..⒇ + # 0x247e: halfwidth smiling face + # 0x247f: halfwidth grinning face + # 0x2480: halfwidth hurt face + # 0x2481: halfwidth frowny face + # 0x2482: halfwidth ⤴ + # 0x2483: halfwidth ⤵ + # 0x2484: halfwidth ZZ ligature + # 0x2485: superscript e + # 0x2486: PK ligature + # 0x2487: MN ligature +} def LittleEndianBitStruct(*args): @@ -488,8 +542,9 @@ def LittleEndianBitStruct(*args): ) class PokemonStringAdapter(Adapter): - u"""Adapter that encodes/decodes Pokémon-formatted text stored in a regular - String struct. + u"""Base adapter for names + + Encodes/decodes Pokémon-formatted text stored in a regular String struct. """ def _decode(self, obj, context): decoded_text = obj.decode('utf16') @@ -499,14 +554,26 @@ def _decode(self, obj, context): decoded_text = decoded_text[0:decoded_text.index(u'\uffff')] # XXX save "trash bytes" somewhere..? - return decoded_text.translate(character_table) + return decoded_text.translate(self.character_table) def _encode(self, obj, context): - #padded_text = (obj + u'\xffff' + '\x00' * 12) - padded_text = obj - decoded_text = padded_text.translate(inverse_character_table) + padded_text = (obj + u'\uffff' + '\x00' * 20) + decoded_text = padded_text.translate(self.inverse_character_table) return decoded_text.encode('utf16') + +def make_pokemon_string_adapter(table, generation): + class _SpecificAdapter(PokemonStringAdapter): + character_table = table + inverse_character_table = dict((ord(v), k) for k, v in + table.iteritems()) + _SpecificAdapter.__name__ = 'PokemonStringAdapterGen%s' % generation + return _SpecificAdapter + +PokemonStringAdapterGen4 = make_pokemon_string_adapter(character_table_gen4, 4) +PokemonStringAdapterGen5 = make_pokemon_string_adapter(character_table_gen5, 5) + + class DateAdapter(Adapter): """Converts between a three-byte string and a Python date. @@ -603,6 +670,11 @@ def make_pokemon_struct(generation): 5: Flag('hidden_ability'), }[generation] + PokemonStringAdapter = { + 4: PokemonStringAdapterGen4, + 5: PokemonStringAdapterGen5, + }[generation] + return Struct('pokemon_struct', # Header ULInt32('personality'), # XXX aughgh http://bulbapedia.bulbagarden.net/wiki/Personality From f2670d2cb606e5b59e0046b429c77b6e08d104f5 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 17 Jun 2012 21:47:12 +0200 Subject: [PATCH 04/23] Enable the PokemonStringAdapter to encode values back into a blob --- pokedex/struct/_pokemon_struct.py | 36 +++++++++++++++++++++++++------ 1 file changed, 29 insertions(+), 7 deletions(-) diff --git a/pokedex/struct/_pokemon_struct.py b/pokedex/struct/_pokemon_struct.py index 40176e5ea..755683ce0 100644 --- a/pokedex/struct/_pokemon_struct.py +++ b/pokedex/struct/_pokemon_struct.py @@ -541,25 +541,47 @@ def LittleEndianBitStruct(*args): resizer=lambda _: _, ) + +class _String(unicode): + pass + + class PokemonStringAdapter(Adapter): u"""Base adapter for names Encodes/decodes Pokémon-formatted text stored in a regular String struct. + + Returns an unicode subclass that has an ``original`` attribute with the + original unencoded value, complete with trash bytes. + On write, if the ``original`` is found, it is written with no regard to the + string value. + This ensures the trash bytes get written back untouched if the string is + unchanged. """ + def __init__(self, field, length): + super(PokemonStringAdapter, self).__init__(field) + self.length = length + def _decode(self, obj, context): decoded_text = obj.decode('utf16') # Real string ends at the \uffff character if u'\uffff' in decoded_text: decoded_text = decoded_text[0:decoded_text.index(u'\uffff')] - # XXX save "trash bytes" somewhere..? - return decoded_text.translate(self.character_table) + result = _String(decoded_text.translate(self.character_table)) + result.original = obj # save original with "trash bytes" + return result def _encode(self, obj, context): - padded_text = (obj + u'\uffff' + '\x00' * 20) + try: + return obj.original + except AttributeError: + pass + length = self.length + padded_text = (obj + u'\uffff' + '\x00' * length) decoded_text = padded_text.translate(self.inverse_character_table) - return decoded_text.encode('utf16') + return decoded_text.encode('utf16')[:length] def make_pokemon_string_adapter(table, generation): @@ -640,7 +662,7 @@ def _encode(self, obj, context): try: forms = self.pokemon_forms[ context['national_id'] ] except KeyError: - return None + return 0 return forms.index(obj) << 3 @@ -830,7 +852,7 @@ def make_pokemon_struct(generation): ULInt16('pt_met_location_id'), # Block C - PokemonStringAdapter(String('nickname', 22)), + PokemonStringAdapter(String('nickname', 22), 22), Padding(1), Enum(ULInt8('original_version'), _unset = 0, @@ -872,7 +894,7 @@ def make_pokemon_struct(generation): Padding(4), # Block D - PokemonStringAdapter(String('original_trainer_name', 16)), + PokemonStringAdapter(String('original_trainer_name', 16), 16), DateAdapter(String('date_egg_received', 3)), DateAdapter(String('date_met', 3)), ULInt16('dp_egg_location_id'), From ab3d4050cfe5f67667faab61d1440fe38b2f42f6 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 24 Jun 2012 23:54:00 +0200 Subject: [PATCH 05/23] =?UTF-8?q?Export=20pok=C3=A9mon=20structures=20as?= =?UTF-8?q?=20dictionaries?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pokedex/struct/__init__.py | 232 +++++++++++++++++++++++++++++- pokedex/struct/_pokemon_struct.py | 8 +- 2 files changed, 231 insertions(+), 9 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index d119e8231..07182bf81 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 @@ def __init__(self, blob=None, encrypted=False, session=None): if session: self.use_database_session(session) + else: + self._session = None @property def as_struct(self): @@ -91,6 +115,107 @@ def as_encrypted(self): # 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 @@ def is_shiny(self): 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 @@ def use_database_session(self, session): 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 @@ def move_pp(self): 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 @@ def reciprocal_crypt(cls, words): 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 @@ def shiny_leaves(self): 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 @@ def use_database_session(self, session): 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 755683ce0..f9590d387 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), From 4dcff2e343d63f9c433516e7e4d9b7695505bf3c Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sat, 13 Oct 2012 16:16:02 +0200 Subject: [PATCH 06/23] Make SaveFilePokemon mutable, import/update from dict --- pokedex/struct/__init__.py | 702 +++++++++++++++++++++--------- pokedex/struct/_pokemon_struct.py | 87 ++-- 2 files changed, 531 insertions(+), 258 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index 07182bf81..f10dc2115 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -9,11 +9,14 @@ """ import struct +import base64 +import datetime +import contextlib -from pokedex.db import tables +from pokedex.db import tables, util from pokedex.formulae import calculated_hp, calculated_stat from pokedex.compatibility import namedtuple, permutations -from pokedex.struct._pokemon_struct import make_pokemon_struct +from pokedex.struct._pokemon_struct import make_pokemon_struct, pokemon_forms def pokemon_prng(seed): u"""Creates a generator that simulates the main Pokémon PRNG.""" @@ -23,27 +26,113 @@ def pokemon_prng(seed): yield seed >> 16 -def _struct_proxy(name): +def struct_proxy(name, dependent=[]): def getter(self): return getattr(self.structure, name) def setter(self, value): setattr(self.structure, name, value) + for dep in dependent: + delattr(self, dep) + del self.blob return property(getter, setter) -def _struct_frozenset_proxy(name): + +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) + bitstruct = dict.fromkeys(value, True) setattr(self.structure, name, bitstruct) + del self.blob return property(getter, setter) +class cached_property(object): + def __init__(self, getter, setter=None): + self._getter = getter + self._setter = setter + self.cache_setter_value = True + + def setter(self, func): + """With this setter, the value being set is automatically cached + """ + self._setter = func + self.cache_setter_value = True + return self + + def complete_setter(self, func): + """Setter without automatic caching of the set value + """ + self._setter = func + self.cache_setter_value = False + return self + + def __get__(self, instance, owner): + if instance is None: + return self + else: + try: + return instance._cached_properties[self] + except AttributeError: + instance._cached_properties = {} + except KeyError: + pass + result = self._getter(instance) + instance._cached_properties[self] = result + return result + + def __set__(self, instance, value): + if self._setter is None: + raise AttributeError('Cannot set attribute') + else: + self._setter(instance, value) + if self.cache_setter_value: + try: + instance._cached_properties[self] = value + except AttributeError: + instance._cached_properties = {self: value} + del instance.blob + + def __delete__(self, instance): + try: + del instance._cached_properties[self] + except (AttributeError, KeyError): + pass + + +class InstrumentedList(object): + def __init__(self, callback, initial=()): + self.list = list(initial) + self.callback = callback + + def __getitem__(self, index): + return self.list[index] + + def __setitem__(self, index, value): + self.list[index] = value + self.callback() + + def __delitem__(self, index, value): + self.list[index] = value + self.callback() + + def append(self, item): + self.list.append(item) + self.callback() + + def extend(self, extralist): + self.list.extend(extralist) + self.callback() + + def __iter__(self): + return iter(self.list) + + class SaveFilePokemon(object): u"""Base class for an individual Pokémon, from the game's point of view. @@ -52,7 +141,7 @@ class SaveFilePokemon(object): """ Stat = namedtuple('Stat', ['stat', 'base', 'gene', 'exp', 'calc']) - def __init__(self, blob=None, encrypted=False, session=None): + def __init__(self, blob=None, dict_=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 @@ -86,14 +175,16 @@ def __init__(self, blob=None, encrypted=False, session=None): # Already decrypted self.blob = blob - self.structure = self.pokemon_struct.parse(self.blob) else: - self.structure = self.pokemon_struct.parse('\0' * (32 * 4 + 8)) + self.blob = '\0' * (32 * 4 + 8) if session: - self.use_database_session(session) + self.session = session else: - self._session = None + self.session = None + + if dict_: + self.update(dict_) @property def as_struct(self): @@ -115,27 +206,32 @@ def as_encrypted(self): # Stuff back into a string, and done return struct.pack(struct_def, *shuffled) - def export(self): + def export_dict(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), ) + ability = self.ability + if ability: + result['ability'] = dict(id=st.ability_id, name=ability.name) - result['original trainer'] = dict( + 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 (trainer['id'] or trainer['secret'] or + trainer['name'].strip('\0') or trainer['gender'] != 'male'): + result['oiginal trainer'] = trainer if self.form != self.species.default_form: - result['form'] = dict(id=self.form.id, name=self.form.form_name) + result['form'] = dict(id=st.form_id, name=self.form.form_name) if self.held_item: - result['item'] = dict(id=self.item.id, name=self.item.name) + result['item'] = dict(id=st.held_item_id, name=self.item.name) if self.exp: result['exp'] = self.exp if self.happiness: @@ -146,15 +242,17 @@ def export(self): 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': + if self.encounter_type and self.encounter_type != 'special': 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, + result['egg location'] = dict( + id=st.pt_egg_location_id or st.dp_egg_location_id, name=self.egg_location.name) if self.met_location: - result['met location'] = dict(id=self.egg_location.id, + result['met location'] = dict( + id=st.pt_met_location_id or st.dp_met_location_id, name=self.met_location.name) if self.date_egg_received: result['egg received'] = self.date_egg_received.isoformat() @@ -163,13 +261,13 @@ def export(self): if self.pokerus: result['pokerus data'] = self.pokerus if self.pokeball: - result['pokeball'] = dict(id=self.pokeball.id, + result['pokeball'] = dict(id=st.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_nicknamed: + result['nicknamed'] = True if self.is_egg: result['is egg'] = True if self.fateful_encounter: @@ -195,12 +293,14 @@ def export(self): 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 pokemon_stat in self.pokemon.stats: + stat_identifier = pokemon_stat.stat.identifier + st_stat_identifier = stat_identifier.replace('-', '_') + dct_stat_identifier = stat_identifier.replace('-', ' ') + if st['iv_' + st_stat_identifier]: + genes[dct_stat_identifier] = st['iv_' + st_stat_identifier] + if st['effort_' + st_stat_identifier]: + effort[dct_stat_identifier] = st['effort_' + st_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] @@ -209,13 +309,124 @@ def export(self): if genes: result['genes'] = genes if contest_stats: - result['contest_stats'] = contest_stats + result['contest stats'] = contest_stats ribbons = list(self.ribbons) if ribbons: result['ribbons'] = ribbons return result + def update(self, dct, **kwargs): + """Updates the pokemon from a YAML/JSON-compatible dict + + Dicts that don't specify all the data are allowed. They update the + structure with the information they contain. + + Keyword arguments with single keys are allowed. The semantics are + similar to dict.update. + + Unlike setting properties directly, the this method tries more to keep + the result sensible, e.g. when species is updated, it can switch + to/from genderless. + """ + st = self.structure + session = self.session + dct.update(kwargs) + reset_form = False + if 'ability' in dct: + st.ability_id = dct['ability']['id'] + del self.ability + if 'form' in dct: + st.alternate_form = dct['form'] + reset_form = True + if 'species' in dct: + st.national_id = dct['species']['id'] + if 'form' not in dct: + st.alternate_form = 0 + reset_form = True + if reset_form: + del self.form + if not self.is_nicknamed: + self.nickname = self.species.name + self.is_nicknamed = False + if self.species.gender_rate == -1: + self.gender = 'genderless' + elif self.gender == 'genderless': + # make id=0 the default, sorry if it looks sexist + self.gender = 'male' + if 'held item' in dct: + st.item_id = dct['held item']['id'] + del self.item + if 'pokeball' in dct: + st.dppt_pokeball = dct['pokeball']['id'] + del self.pokeball + def _load_values(source, **values): + for attrname, key in values.iteritems(): + try: + value = source[key] + except KeyError: + pass + else: + setattr(self, attrname, value) + if 'oiginal trainer' in dct: + _load_values(dct['oiginal trainer'], + original_trainer_id='id', + original_trainer_secret_id='secret', + original_trainer_name='name', + original_trainer_gender='gender', + ) + n = self.is_nicknamed + _load_values(dct, + exp='exp', + happiness='happiness', + markings='markings', + original_country='original country', + original_version='original version', + encounter_type='encounter type', + nickname='nickname', + pokerus='pokerus data', + met_at_level='met at level', + is_nicknamed='nicknamed', + is_egg='is egg', + fateful_encounter='fateful encounter', + gender='gender', + ) + self.is_nicknamed = n + if 'egg location' in dct: + st.pt_egg_location_id = dct['egg location']['id'] + del self.egg_location + if 'met location' in dct: + st.pt_met_location_id = dct['met location']['id'] + del self.met_location + if 'date met' in dct: + self.date_met = datetime.datetime.strptime( + dct['date met'], '%Y-%m-%d').date() + if 'date egg received' in dct: + self.egg_received = datetime.datetime.strptime( + dct['date egg received'], '%Y-%m-%d').date() + if 'moves' in dct: + pp_reset_indices = [] + for i, movedict in enumerate(dct['moves']): + setattr(st, 'move{0}_id'.format(i + 1), movedict['id']) + if 'pp' in movedict: + setattr(st, 'move{0}_pp'.format(i + 1), movedict['pp']) + else: + pp_sets.append(i) + for i in range(i + 1, 4): + # Reset the rest of the moves + setattr(st, 'move{0}_id'.format(i + 1), 0) + setattr(st, 'move{0}_pp'.format(i + 1), 0) + del self.moves + del self.move_pp + for i in pp_reset_indices: + # Set default PP here, when the moves dict is regenerated + setattr(st, 'move{0}_pp'.format(i + 1), self.moves[i].pp) + for key, prefix in (('genes', 'iv'), ('effort', 'ev'), + ('contest stats', 'contest')): + for name, value in dct.get(key, {}).items(): + st['{}_{}'.format(prefix, name.replace(' ', '_'))] = value + return self + ### Delicious data @property def is_shiny(self): @@ -237,82 +448,105 @@ def use_database_session(self, session): database stuff. Gotta call this (or give session to `__init__`) before you use the database properties like `species`, etc. """ - self._session = session + if self.session and self.session is not session: + raise ValueError('Re-setting a session is not supported') + self.session = session + + @cached_property + def stats(self): + stats = [] + for pokemon_stat in self.pokemon.stats: + stat_identifier = pokemon_stat.stat.identifier.replace('-', '_') + gene = st['iv_' + stat_identifier] + exp = st['effort_' + stat_identifier] + if pokemon_stat.stat.identifier == 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, + ), + ) + + stats.append(stat_tup) + return tuple(stats) + + @property + def alternate_form(self): st = self.structure + forms = pokemon_forms.get(st.national_id) + if forms: + return forms[st.alternate_form_id] + else: + return None - 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(form_identifier=st.alternate_form) \ - .one() + @alternate_form.setter + def alternate_form(self, alternate_form): + st = self.structure + forms = pokemon_forms.get(st.national_id) + if forms: + st.alternate_form_id = forms.index(alternate_form) else: - self._pokemon = self._pokemon_form = None - self._ability = self._session.query(tables.Ability).get(st.ability_id) - - 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) \ - .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 - - if self._pokemon: - self._stats = [] - for pokemon_stat in self._pokemon.stats: - stat_identifier = pokemon_stat.stat.identifier.replace('-', '_') - gene = st['iv_' + stat_identifier] - exp = st['effort_' + stat_identifier] - - if pokemon_stat.stat.identifier == 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, - ), - ) + st.alternate_form_id = 0 + del self.form - self._stats.append(stat_tup) + @property + def species(self): + if self.form: + return self.form.species else: - self._stats = [0] * 6 + return None + @species.setter + def species(self, species): + self.form = species.default_form - 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) + @property + def pokemon(self): + if self.form: + return self.form.pokemon + else: + return None - self._moves = [moves_dict.get(move_id, None) for move_id in move_ids] + @pokemon.setter + def pokemon(self, pokemon): + self.form = pokemon.default_form + + @cached_property + def form(self): + st = self.structure + session = self.session + if st.national_id: + pokemon = session.query(tables.Pokemon).get(st.national_id) + return session.query(tables.PokemonForm) \ + .with_parent(pokemon) \ + .filter_by(form_identifier=self.alternate_form) \ + .one() + else: + return None + + @form.setter + def form(self, form): + self.structure.national_id = form.species.id + self.structure.alternate_form = form.form_identifier + del self.species + del self.pokemon + del self.ability + self._reset() + @property + def pokeball(self): + st = self.structure if st.hgss_pokeball >= 17: pokeball_id = st.hgss_pokeball - 17 + 492 elif st.dppt_pokeball: @@ -320,88 +554,125 @@ def use_database_session(self, session): else: pokeball_id = None if pokeball_id: - self._pokeball = session.query(tables.ItemGameIndex) \ - .filter_by(generation_id = 4, game_index = pokeball_id).one().item + self._pokeball = self.session.query(tables.ItemGameIndex) \ + .filter_by(generation_id = self.generation_id, + game_index = pokeball_id).one().item + @cached_property + def egg_location(self): + st = self.structure 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 = 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 = self.generation_id, game_index = met_loc_id).one().location + return self.session.query(tables.LocationGameIndex) \ + .filter_by(generation_id=4, + game_index = egg_loc_id).one().location else: - self._met_location = None - - @property - def species(self): - return self._pokemon_form.species - - @property - def pokemon(self): - return self._pokemon_form.pokemon - - @property - def form(self): - return self._pokemon_form - - @property - def species_form(self): - return self._pokemon_form + return None - @property - def pokeball(self): - return self._pokeball - - @property - def egg_location(self): - return self._egg_location - - @property + @cached_property def met_location(self): - return self._met_location + st = self.structure + met_loc_id = st.pt_met_location_id or st.dp_met_location_id + if met_loc_id: + return self.session.query(tables.LocationGameIndex) \ + .filter_by(generation_id=4, + game_index=met_loc_id).one().location + else: + return None @property def level(self): - return self._experience_rung.level + return self.experience_rung.level + + @cached_property + def experience_rung(self): + growth_rate = self.species.growth_rate + return (session.query(tables.Experience) + .filter(tables.Experience.growth_rate == growth_rate) + .filter(tables.Experience.experience <= self.exp) + .order_by(tables.Experience.level.desc()) + [0]) + + @cached_property + def next_experience_rung(self): + level = self.level + if level < 100: + return (session.query(tables.Experience) + .filter(tables.Experience.growth_rate == growth_rate) + .filter(tables.Experience.level == level + 1) + .one()) + else: + return None @property def exp_to_next(self): - if self._next_experience_rung: - return self._next_experience_rung.experience - self.structure.exp + if self.next_experience_rung: + return self.next_experience_rung.experience - self.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) + if self.next_experience_rung: + rung = self.experience_rung + return (1.0 * + (self.exp - rung.experience) / + (self.next_experience_rung.experience - rung.experience)) else: return 0.0 - @property + @cached_property def ability(self): - return self._ability + return self.session.query(tables.Ability).get(self.structure.ability_id) - @property + @ability.setter + def ability(self, ability): + self.structure.ability_id = ability.id + + @cached_property def held_item(self): - return self._held_item + held_item_id = self.structure.held_item_id + if held_item_id: + return session.query(tables.ItemGameIndex) \ + .filter_by(generation_id=self.generation_id, + game_index=held_item_id) \ + .one().item + + @cached_property + def moves(self): + 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) - @property - def stats(self): - return self._stats + def callback(): + def get(index): + try: + return result[x].id + except AttributeError: + return 0 + self.structure.move1_id = get(0) + self.structure.move2_id = get(1) + self.structure.move3_id = get(2) + self.structure.move4_id = get(3) + self._reset() + + result = InstrumentedList( + callback, + [moves_dict.get(move_id, None) for move_id in move_ids]) - @property - def moves(self): - return self._moves + return result - @property + @moves.complete_setter + def moves(self, new_moves): + self.moves[:] = new_moves + + @cached_property def move_pp(self): return ( self.structure.move1_pp, @@ -410,36 +681,42 @@ def move_pp(self): self.structure.move4_pp, ) - @property + @move_pp.complete_setter + def move_pp(self, new_pps): + self.move_pp[:] = new_pps + + @cached_property def markings(self): return frozenset(k for k, v in self.structure.markings.items() if v) - @markings.setter + @markings.complete_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') + del self.markings + + 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', + dependent=['experience_rung', 'next_experience_rung']) + 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): @@ -456,12 +733,14 @@ def nickname(self): @nickname.setter def nickname(self, value): self.structure.nickname = value - self.structure.is_nicknamed = True + self.is_nicknamed = True + del self.blob @nickname.deleter def nickname(self, value): self.structure.nickname = '' - self.structure.is_nicknamed = False + self.is_nicknamed = False + del self.blob ### Utility methods @@ -519,31 +798,38 @@ def reciprocal_crypt(cls, words): 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. + @cached_property + def blob(self): + """Update the blob and checksum with modified structure """ - 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 + blob = self.pokemon_struct.build(self.structure) + self.structure = self.pokemon_struct.parse(blob) + checksum = sum(struct.unpack('H' * 0x40, 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) + blob = blob[:6] + struct.pack('H', checksum) + blob[8:] + return blob + + @blob.setter + def blob(self, blob): + self.structure = self.pokemon_struct.parse(blob) class SaveFilePokemonGen4(SaveFilePokemon): generation_id = 4 pokemon_struct = make_pokemon_struct(generation=generation_id) - def export(self): - result = super(SaveFilePokemonGen5, self).export() + def export_dict(self): + result = super(SaveFilePokemonGen5, self).export_dict() if any(self.shiny_leaves): result['shiny leaves'] = self.shiny_leaves return result + def update(self, dct, **kwargs): + dct.update(kwargs) + if 'shiny leaves' in dct: + self.shiny_leaves = dct['shiny leaves'] + super(SaveFilePokemonGen4, self).update(dct) + @property def shiny_leaves(self): return ( @@ -565,37 +851,41 @@ def shiny_leaves(self, new_values): self.structure.shining_leaves.leaf5, self.structure.shining_leaves.crown, ) = new_values - self._reset() + del self.blob 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() - - def export(self): - result = super(SaveFilePokemonGen5, self).export() - result['nature'] = dict(id=self.nature.id, name=self.nature.name) + def export_dict(self): + result = super(SaveFilePokemonGen5, self).export_dict() + if self.nature: + result['nature'] = dict( + id=self.structure.nature_id, name=self.nature.name) return result + def update(self, dct, **kwargs): + dct.update(kwargs) + if 'nature' in dct: + self.structure.nature_id = dct['nature']['id'] + del self.nature + super(SaveFilePokemonGen5, self).update(dct) + # XXX: Ability setter must set hidden ability flag - @property + @cached_property def nature(self): - return self._nature + st = self.structure + if st.nature_id: + return (self.session.query(tables.Nature) + .filter_by(game_index = st.nature_id).one()) + else: + return None @nature.setter def nature(self, new_nature): self.structure.nature_id = int(new_nature.game_index) - self._reset() save_file_pokemon_classes = { diff --git a/pokedex/struct/_pokemon_struct.py b/pokedex/struct/_pokemon_struct.py index f9590d387..256ec9fba 100644 --- a/pokedex/struct/_pokemon_struct.py +++ b/pokedex/struct/_pokemon_struct.py @@ -17,6 +17,38 @@ # - higher-level validation; see XXXes below # - personality indirectly influences IVs due to PRNG use +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', '???', + ], +} + # The entire gen 4 character table: character_table_gen4 = { 0x0002: u'ぁ', @@ -616,57 +648,6 @@ def _encode(self, obj, context): 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 0 - - return forms.index(obj) << 3 - - # Docs: http://projectpokemon.org/wiki/Pokemon_NDS_Structure # http://projectpokemon.org/wiki/Pokemon_Black/White_NDS_Structure # http://projectpokemon.org/forums/showthread.php?11474-Hex-Values-and-Trashbytes-in-B-W#post93598 @@ -729,6 +710,7 @@ def make_pokemon_struct(generation): de=5, es=7, kr=8, + unknown_193=193, ), # XXX sum cannot surpass 510 @@ -837,7 +819,7 @@ def make_pokemon_struct(generation): Flag('cool_ribbon'), ), Embed(EmbeddedBitStruct( - PokemonFormAdapter(BitField('alternate_form', 5)), + BitField('alternate_form_id', 5), Enum(BitField('gender', 2), genderless = 2, male = 0, @@ -867,6 +849,7 @@ def make_pokemon_struct(generation): pearl = 11, platinum = 12, orre = 15, + event_20 = 20, ), LittleEndianBitStruct('sinnoh_contest_ribbons', Padding(12), From 93aa4ca60ebc3d1e2c0974784df6279d2389ea4b Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sat, 13 Oct 2012 22:31:18 +0200 Subject: [PATCH 07/23] Add the pkm subcommand to pokedex --- pokedex/main.py | 84 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) diff --git a/pokedex/main.py b/pokedex/main.py index fe4557730..c1319dbab 100644 --- a/pokedex/main.py +++ b/pokedex/main.py @@ -2,11 +2,15 @@ from optparse import OptionParser import os import sys +import textwrap +import json +import base64 import pokedex.db import pokedex.db.load import pokedex.db.tables import pokedex.lookup +import pokedex.struct from pokedex import defaults def main(): @@ -260,6 +264,85 @@ def command_lookup(*args): print +def command_pkm(*args): + if args and args[0] == 'encode': + mode = 'encode' + elif args and args[0] == 'decode': + mode = 'decode' + else: + print textwrap.dedent(u""" + Convert binary Pokémon data (aka PKM files) to/from JSON/YAML. + usage: pokedex pkm (encode|decode) [options] ... + + Commands: + encode Convert a JSON or YAML representation of a + Pokémon to the binary format. + decode Convert the binary format to a JSON/YAML + representation. + + Options: + --gen=NUM, -g Generation to use (4 or 5) + --crypt, -c Use encrypted binary format. + --yaml, -y Use YAML as the human-readable format. + Requires the PyYAML library to be installed. + Default is to use JSON. + --base64, -b Use Base64 encoding for the binary format. + By default, on for 'encode' but off for 'decode' + --binary, -B Output raw binary data, do not use Base64 + + If no files are given, reads from standard input. + """).encode(sys.getdefaultencoding(), 'replace') + return + parser = get_parser(verbose=False) + parser.add_option('-g', '--gen', default=5, type=int) + parser.add_option('-c', '--crypt', action='store_true') + parser.add_option('-y', '--yaml', action='store_true') + parser.add_option('-b', '--base64', action='store_true', + default=(mode == 'encode')) + parser.add_option('-B', '--no-base64', action='store_false', dest='base64') + options, files = parser.parse_args(list(args[1:])) + + session = get_session(options) + cls = pokedex.struct.save_file_pokemon_classes[options.gen] + if options.yaml: + import yaml + + if not files: + # Use sys.stdin in place of name, handle specially later + files = [sys.stdin] + + for filename in files: + if filename is sys.stdin: + content = sys.stdin.read() + else: + with open(filename) as f: + content = f.read() + if mode == 'encode': + if options.yaml: + dict_ = yaml.load(content) + else: + dict_ = json.loads(content) + struct = cls(session=session, dict_=dict_) + if options.crypt: + data = struct.as_encrypted + else: + data = struct.as_struct + if options.base64: + print base64.b64encode(data) + else: + sys.stdout.write(data) + else: + if options.base64: + content = base64.b64decode(content) + struct = cls( + blob=content, encrypted=options.crypt, session=session) + dict_ = struct.export_dict() + if options.yaml: + print yaml.safe_dump(dict_, explicit_start=True), + else: + print json.dumps(dict_), + + def command_help(): print u"""pokedex -- a command-line Pokédex interface usage: pokedex {command} [options...] @@ -269,6 +352,7 @@ def command_help(): Commands: help Displays this message. lookup [thing] Look up something in the Pokédex. + pkm Binary Pokémon format encoding/decoding. (experimental) System commands: load Load Pokédex data into a database from CSV files. From b0bedfb5de2580766a49d606d563c3f06c8a903d Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sat, 13 Oct 2012 22:43:25 +0200 Subject: [PATCH 08/23] Allow some "enums" in the structure to have unknown int values --- pokedex/struct/_pokemon_struct.py | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/pokedex/struct/_pokemon_struct.py b/pokedex/struct/_pokemon_struct.py index 256ec9fba..555ed6753 100644 --- a/pokedex/struct/_pokemon_struct.py +++ b/pokedex/struct/_pokemon_struct.py @@ -648,6 +648,21 @@ def _encode(self, obj, context): y, m, d = obj.year - 2000, obj.month, obj.day return ''.join(chr(n) for n in (y, m, d)) +class LeakyEnum(Adapter): + """An Enum that allows unknown values""" + def __init__(self, sub, **values): + super(LeakyEnum, self).__init__(sub) + self.values = values + self.inverted_values = dict((v, k) for k, v in values.items()) + assert len(values) == len(self.inverted_values) + + def _encode(self, obj, context): + return self.inverted_values.get(obj, obj) + + def _decode(self, obj, context): + return self.values.get(obj, obj) + + # Docs: http://projectpokemon.org/wiki/Pokemon_NDS_Structure # http://projectpokemon.org/wiki/Pokemon_Black/White_NDS_Structure # http://projectpokemon.org/forums/showthread.php?11474-Hex-Values-and-Trashbytes-in-B-W#post93598 @@ -701,7 +716,7 @@ def make_pokemon_struct(generation): Flag('triangle'), Flag('circle'), ), - Enum(ULInt8('original_country'), + LeakyEnum(ULInt8('original_country'), _unset = 0, jp=1, us=2, @@ -710,7 +725,6 @@ def make_pokemon_struct(generation): de=5, es=7, kr=8, - unknown_193=193, ), # XXX sum cannot surpass 510 @@ -836,7 +850,7 @@ def make_pokemon_struct(generation): # Block C PokemonStringAdapter(String('nickname', 22), 22), Padding(1), - Enum(ULInt8('original_version'), + LeakyEnum(ULInt8('original_version'), _unset = 0, sapphire = 1, ruby = 2, @@ -849,7 +863,6 @@ def make_pokemon_struct(generation): pearl = 11, platinum = 12, orre = 15, - event_20 = 20, ), LittleEndianBitStruct('sinnoh_contest_ribbons', Padding(12), @@ -891,7 +904,7 @@ def make_pokemon_struct(generation): ), BitField('met_at_level', 7), ), - Enum(ULInt8('encounter_type'), + LeakyEnum(ULInt8('encounter_type'), special = 0, # egg; pal park; event; honey tree; shaymin grass = 2, # or darkrai dialga_palkia = 4, From 8590c51c7c6e26e6cd82de5f1378114b3c092e89 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sat, 13 Oct 2012 22:01:49 +0200 Subject: [PATCH 09/23] Saving string trash bytes, bug fixes --- pokedex/struct/__init__.py | 216 +++++++++++++++++++----------- pokedex/struct/_pokemon_struct.py | 24 ++-- 2 files changed, 152 insertions(+), 88 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index f10dc2115..832b22f7c 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -13,10 +13,13 @@ import datetime import contextlib +import sqlalchemy.orm.exc + from pokedex.db import tables, util from pokedex.formulae import calculated_hp, calculated_stat from pokedex.compatibility import namedtuple, permutations -from pokedex.struct._pokemon_struct import make_pokemon_struct, pokemon_forms +from pokedex.struct._pokemon_struct import (make_pokemon_struct, pokemon_forms, + StringWithOriginal) def pokemon_prng(seed): u"""Creates a generator that simulates the main Pokémon PRNG.""" @@ -45,8 +48,9 @@ def getter(self): return frozenset(k for k, v in bitstruct.items() if v) def setter(self, value): - bitstruct = dict.fromkeys(value, True) - setattr(self.structure, name, bitstruct) + struct = getattr(self.structure, name) + for key in struct: + setattr(struct, key, key in value) del self.blob return property(getter, setter) @@ -211,6 +215,13 @@ def export_dict(self): """ st = self.structure + def save_trash(result, name, string): + trash = getattr(string, 'original', None) + if trash: + expected = (string + u'\uffff').encode('utf-16LE') + if trash.rstrip('\0') != expected: + result[name] = base64.b64encode(trash) + result = dict( species=dict(id=self.species.id, name=self.species.name), ) @@ -224,6 +235,7 @@ def export_dict(self): name=unicode(self.original_trainer_name), gender=self.original_trainer_gender ) + save_trash(trainer, 'name trash', self.original_trainer_name) if (trainer['id'] or trainer['secret'] or trainer['name'].strip('\0') or trainer['gender'] != 'male'): result['oiginal trainer'] = trainer @@ -231,7 +243,7 @@ def export_dict(self): if self.form != self.species.default_form: result['form'] = dict(id=st.form_id, name=self.form.form_name) if self.held_item: - result['item'] = dict(id=st.held_item_id, name=self.item.name) + result['held item'] = dict(id=st.held_item_id, name=self.held_item.name) if self.exp: result['exp'] = self.exp if self.happiness: @@ -246,14 +258,25 @@ def export_dict(self): result['encounter type'] = self.encounter_type if self.nickname: result['nickname'] = unicode(self.nickname) + save_trash(result, 'nickname trash', self.nickname) if self.egg_location: result['egg location'] = dict( id=st.pt_egg_location_id or st.dp_egg_location_id, name=self.egg_location.name) + elif st.pt_egg_location_id or st.dp_egg_location_id: + result['egg location'] = dict( + id=st.pt_egg_location_id or st.dp_egg_location_id) + if st.dp_egg_location_id: + result['egg location slot'] = 'dp' if self.met_location: result['met location'] = dict( id=st.pt_met_location_id or st.dp_met_location_id, name=self.met_location.name) + elif st.pt_met_location_id or st.dp_met_location_id: + result['egg location'] = dict( + id=st.pt_met_location_id or st.dp_met_location_id) + if st.dp_met_location_id: + result['met location slot'] = 'dp' if self.date_egg_received: result['egg received'] = self.date_egg_received.isoformat() if self.date_met: @@ -261,7 +284,8 @@ def export_dict(self): if self.pokerus: result['pokerus data'] = self.pokerus if self.pokeball: - result['pokeball'] = dict(id=st.pokeball_id, + result['pokeball'] = dict( + id=st.dppt_pokeball or st.hgss_pokeball, name=self.pokeball.name) if self.met_at_level: result['met at level'] = self.met_at_level @@ -272,21 +296,24 @@ def export_dict(self): result['is egg'] = True if self.fateful_encounter: result['fateful encounter'] = True + if self.personality: + result['personality'] = self.personality if self.gender != 'genderless': result['gender'] = self.gender + if st.hidden_ability: + result['has hidden ability'] = st.hidden_ability + 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 + move['pp'] = st['move%s_pp' % i] pp_up = st['move%s_pp_ups' % i] if pp_up: - move['pp_up'] = pp_up + move['pp ups'] = pp_up if move: moves.append(move) @@ -297,21 +324,18 @@ def export_dict(self): stat_identifier = pokemon_stat.stat.identifier st_stat_identifier = stat_identifier.replace('-', '_') dct_stat_identifier = stat_identifier.replace('-', ' ') - if st['iv_' + st_stat_identifier]: - genes[dct_stat_identifier] = st['iv_' + st_stat_identifier] - if st['effort_' + st_stat_identifier]: - effort[dct_stat_identifier] = st['effort_' + st_stat_identifier] + genes[dct_stat_identifier] = st['iv_' + st_stat_identifier] + effort[dct_stat_identifier] = st['effort_' + st_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: + contest_stats[contest_stat] = st['contest_' + contest_stat] + if any(effort.values()): result['effort'] = effort - if genes: + if any(genes.values()): result['genes'] = genes - if contest_stats: + if any(contest_stats.values()): result['contest stats'] = contest_stats - ribbons = list(self.ribbons) + ribbons = sorted(r.replace('_', ' ') for r in self.ribbons) if ribbons: result['ribbons'] = ribbons return result @@ -332,10 +356,10 @@ def update(self, dct, **kwargs): st = self.structure session = self.session dct.update(kwargs) - reset_form = False if 'ability' in dct: st.ability_id = dct['ability']['id'] del self.ability + reset_form = False if 'form' in dct: st.alternate_form = dct['form'] reset_form = True @@ -355,10 +379,10 @@ def update(self, dct, **kwargs): # make id=0 the default, sorry if it looks sexist self.gender = 'male' if 'held item' in dct: - st.item_id = dct['held item']['id'] - del self.item + st.held_item_id = dct['held item']['id'] + del self.held_item if 'pokeball' in dct: - st.dppt_pokeball = dct['pokeball']['id'] + self.pokeball = self._get_pokeball(dct['pokeball']['id']) del self.pokeball def _load_values(source, **values): for attrname, key in values.iteritems(): @@ -368,14 +392,23 @@ def _load_values(source, **values): pass else: setattr(self, attrname, value) + def load_name(attr_name, dct, string_key, trash_key): + if string_key in dct: + if trash_key in dct: + name = StringWithOriginal(unicode(dct[string_key])) + name.original = base64.b64decode(dct[trash_key]) + setattr(self, attr_name, name) + else: + setattr(self, attr_name, unicode(dct[string_key])) if 'oiginal trainer' in dct: - _load_values(dct['oiginal trainer'], + trainer = dct['oiginal trainer'] + _load_values(trainer, original_trainer_id='id', original_trainer_secret_id='secret', - original_trainer_name='name', original_trainer_gender='gender', ) - n = self.is_nicknamed + load_name('original_trainer_name', trainer, 'name', 'name trash') + was_nicknamed = self.is_nicknamed _load_values(dct, exp='exp', happiness='happiness', @@ -383,45 +416,60 @@ def _load_values(source, **values): original_country='original country', original_version='original version', encounter_type='encounter type', - nickname='nickname', pokerus='pokerus data', met_at_level='met at level', - is_nicknamed='nicknamed', is_egg='is egg', fateful_encounter='fateful encounter', gender='gender', + personality='personality', + hidden_ability='has hidden ability', + ) + load_name('nickname', dct, 'nickname', 'nickname trash') + self.is_nicknamed = was_nicknamed + _load_values(dct, + is_nicknamed='nicknamed', ) - self.is_nicknamed = n - if 'egg location' in dct: - st.pt_egg_location_id = dct['egg location']['id'] - del self.egg_location - if 'met location' in dct: - st.pt_met_location_id = dct['met location']['id'] - del self.met_location + for loc_type in 'egg', 'met': + key = '{0} location'.format(loc_type) + if key in dct: + dp_attr = 'dp_{0}_location_id'.format(loc_type) + pt_attr = 'pt_{0}_location_id'.format(loc_type) + if dct.get('{0} location slot'.format(loc_type)) == 'dp': + attr = dp_attr + other_attr = pt_attr + else: + attr = pt_attr + other_attr = dp_attr + st[attr] = dct[key]['id'] + st[other_attr] = 0 + delattr(self, '{0}_location'.format(loc_type)) if 'date met' in dct: self.date_met = datetime.datetime.strptime( dct['date met'], '%Y-%m-%d').date() - if 'date egg received' in dct: - self.egg_received = datetime.datetime.strptime( - dct['date egg received'], '%Y-%m-%d').date() + if 'egg received' in dct: + self.date_egg_received = datetime.datetime.strptime( + dct['egg received'], '%Y-%m-%d').date() if 'moves' in dct: pp_reset_indices = [] for i, movedict in enumerate(dct['moves']): - setattr(st, 'move{0}_id'.format(i + 1), movedict['id']) + st['move{0}_id'.format(i + 1)] = movedict['id'] if 'pp' in movedict: - setattr(st, 'move{0}_pp'.format(i + 1), movedict['pp']) + st['move{0}_pp'.format(i + 1)] = movedict['pp'] else: - pp_sets.append(i) + pp_reset_indices.append(i) + if 'pp ups' in movedict: + st['move{0}_pp_ups'.format(i + 1)] = movedict['pp ups'] for i in range(i + 1, 4): # Reset the rest of the moves - setattr(st, 'move{0}_id'.format(i + 1), 0) - setattr(st, 'move{0}_pp'.format(i + 1), 0) + st['move{0}_id'.format(i + 1)] = 0 + st['move{0}_pp'.format(i + 1)] = 0 + st['move{0}_pp_up'.format(i + 1)] = 0 del self.moves del self.move_pp for i in pp_reset_indices: # Set default PP here, when the moves dict is regenerated - setattr(st, 'move{0}_pp'.format(i + 1), self.moves[i].pp) - for key, prefix in (('genes', 'iv'), ('effort', 'ev'), + st['move{0}_pp'.format(i + 1)] = self.moves[i].pp + for key, prefix in (('genes', 'iv'), ('effort', 'effort'), ('contest stats', 'contest')): for name, value in dct.get(key, {}).items(): st['{}_{}'.format(prefix, name.replace(' ', '_'))] = value @@ -528,10 +576,13 @@ def form(self): session = self.session if st.national_id: pokemon = session.query(tables.Pokemon).get(st.national_id) - return session.query(tables.PokemonForm) \ - .with_parent(pokemon) \ - .filter_by(form_identifier=self.alternate_form) \ - .one() + if self.alternate_form: + return session.query(tables.PokemonForm) \ + .with_parent(pokemon) \ + .filter_by(form_identifier=self.alternate_form) \ + .one() + else: + return pokemon.default_form else: return None @@ -541,10 +592,9 @@ def form(self, form): self.structure.alternate_form = form.form_identifier del self.species del self.pokemon - del self.ability self._reset() - @property + @cached_property def pokeball(self): st = self.structure if st.hgss_pokeball >= 17: @@ -552,20 +602,36 @@ def pokeball(self): elif st.dppt_pokeball: pokeball_id = st.dppt_pokeball else: - pokeball_id = None - if pokeball_id: - self._pokeball = self.session.query(tables.ItemGameIndex) \ - .filter_by(generation_id = self.generation_id, - game_index = pokeball_id).one().item + return None + return self._get_pokeball(pokeball_id) + + def _get_pokeball(self, pokeball_id): + return (self.session.query(tables.ItemGameIndex) + .filter_by(generation_id=4, game_index = pokeball_id).one().item) + + @pokeball.setter + def pokeball(self, pokeball): + st = self.structure + st.hgss_pokeball = st.dppt_pokeball = 0 + if pokeball: + pokeball_id = pokeball.id + boundary = 492 - 17 + if pokeball_id >= boundary: + st.hgss_pokeball = pokeball_id - boundary + else: + st.dppt_pokeball = pokeball_id @cached_property def egg_location(self): st = self.structure egg_loc_id = st.pt_egg_location_id or st.dp_egg_location_id if egg_loc_id: - return self.session.query(tables.LocationGameIndex) \ - .filter_by(generation_id=4, - game_index = egg_loc_id).one().location + try: + return self.session.query(tables.LocationGameIndex) \ + .filter_by(generation_id=4, + game_index = egg_loc_id).one().location + except sqlalchemy.orm.exc.NoResultFound: + return None else: return None @@ -574,9 +640,12 @@ def met_location(self): st = self.structure met_loc_id = st.pt_met_location_id or st.dp_met_location_id if met_loc_id: - return self.session.query(tables.LocationGameIndex) \ - .filter_by(generation_id=4, - game_index=met_loc_id).one().location + try: + return self.session.query(tables.LocationGameIndex) \ + .filter_by(generation_id=4, + game_index=met_loc_id).one().location + except sqlalchemy.orm.exc.NoResultFound: + return None else: return None @@ -633,7 +702,7 @@ def ability(self, ability): def held_item(self): held_item_id = self.structure.held_item_id if held_item_id: - return session.query(tables.ItemGameIndex) \ + return self.session.query(tables.ItemGameIndex) \ .filter_by(generation_id=self.generation_id, game_index=held_item_id) \ .one().item @@ -685,15 +754,6 @@ def move_pp(self): def move_pp(self, new_pps): self.move_pp[:] = new_pps - @cached_property - def markings(self): - return frozenset(k for k, v in self.structure.markings.items() if v) - - @markings.complete_setter - def markings(self, value): - self.structure.markings = dict((k, True) for k in value) - del self.markings - 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') @@ -712,6 +772,7 @@ def markings(self, value): met_at_level = struct_proxy('met_at_level') original_trainer_gender = struct_proxy('original_trainer_gender') encounter_type = struct_proxy('encounter_type') + personality = struct_proxy('personality') markings = struct_frozenset_proxy('markings') sinnoh_ribbons = struct_frozenset_proxy('sinnoh_ribbons') @@ -800,8 +861,6 @@ def reciprocal_crypt(cls, words): @cached_property def blob(self): - """Update the blob and checksum with modified structure - """ blob = self.pokemon_struct.build(self.structure) self.structure = self.pokemon_struct.parse(blob) checksum = sum(struct.unpack('H' * 0x40, blob[8:0x88])) & 0xffff @@ -867,12 +926,11 @@ def export_dict(self): def update(self, dct, **kwargs): dct.update(kwargs) + super(SaveFilePokemonGen5, self).update(dct) if 'nature' in dct: self.structure.nature_id = dct['nature']['id'] - del self.nature - super(SaveFilePokemonGen5, self).update(dct) - - # XXX: Ability setter must set hidden ability flag + if 'has hidden ability' not in dct: + self.hidden_ability = (self.pokemon.dream_ability == self.ability) @cached_property def nature(self): @@ -887,6 +945,8 @@ def nature(self): def nature(self, new_nature): self.structure.nature_id = int(new_nature.game_index) + hidden_ability = struct_proxy('hidden_ability') + save_file_pokemon_classes = { 4: SaveFilePokemonGen4, diff --git a/pokedex/struct/_pokemon_struct.py b/pokedex/struct/_pokemon_struct.py index 555ed6753..ebb453311 100644 --- a/pokedex/struct/_pokemon_struct.py +++ b/pokedex/struct/_pokemon_struct.py @@ -574,7 +574,7 @@ def LittleEndianBitStruct(*args): ) -class _String(unicode): +class StringWithOriginal(unicode): pass @@ -601,19 +601,23 @@ def _decode(self, obj, context): if u'\uffff' in decoded_text: decoded_text = decoded_text[0:decoded_text.index(u'\uffff')] - result = _String(decoded_text.translate(self.character_table)) + result = StringWithOriginal( + decoded_text.translate(self.character_table)) result.original = obj # save original with "trash bytes" return result def _encode(self, obj, context): try: - return obj.original + original = obj.original except AttributeError: - pass - length = self.length - padded_text = (obj + u'\uffff' + '\x00' * length) - decoded_text = padded_text.translate(self.inverse_character_table) - return decoded_text.encode('utf16')[:length] + length = self.length + padded_text = (obj + u'\uffff' + '\x00' * length) + decoded_text = padded_text.translate(self.inverse_character_table) + return decoded_text.encode('utf-16LE')[:length] + else: + if self._decode(original, context) != obj: + raise ValueError("String and original don't match") + return original def make_pokemon_string_adapter(table, generation): @@ -657,10 +661,10 @@ def __init__(self, sub, **values): assert len(values) == len(self.inverted_values) def _encode(self, obj, context): - return self.inverted_values.get(obj, obj) + return self.values.get(obj, obj) def _decode(self, obj, context): - return self.values.get(obj, obj) + return self.inverted_values.get(obj, obj) # Docs: http://projectpokemon.org/wiki/Pokemon_NDS_Structure From 14ea7a0adf41530247f54138ccf90fbad330591e Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sat, 13 Oct 2012 23:52:51 +0200 Subject: [PATCH 10/23] Fix frozenset proxy, implement ribbons setter --- pokedex/struct/__init__.py | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index 832b22f7c..9bf6e5ffc 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -47,10 +47,14 @@ def getter(self): bitstruct = getattr(self.structure, name) return frozenset(k for k, v in bitstruct.items() if v) - def setter(self, value): + def setter(self, new_set): + new_set = set(new_set) struct = getattr(self.structure, name) for key in struct: - setattr(struct, key, key in value) + setattr(struct, key, key in new_set) + new_set.discard(key) + if new_set: + raise ValueError('Unknown values: {0}'.format(', '.join(ribbons))) del self.blob return property(getter, setter) @@ -449,6 +453,8 @@ def load_name(attr_name, dct, string_key, trash_key): if 'egg received' in dct: self.date_egg_received = datetime.datetime.strptime( dct['egg received'], '%Y-%m-%d').date() + if 'ribbons' in dct: + self.ribbons = (r.replace(' ', '_') for r in dct['ribbons']) if 'moves' in dct: pp_reset_indices = [] for i, movedict in enumerate(dct['moves']): @@ -785,7 +791,19 @@ def ribbons(self): self.sinnoh_ribbons | self.hoenn_ribbons | self.sinnoh_contest_ribbons) - # XXX: ribbons setter + + @ribbons.setter + def ribbons(self, ribbons): + ribbons = set(ribbons) + for ribbonset_name in ( + 'sinnoh_ribbons', 'hoenn_ribbons', 'sinnoh_contest_ribbons'): + ribbonset = self.structure[ribbonset_name] + print ribbonset + for ribbon_name in ribbonset: + ribbonset[ribbon_name] = (ribbon_name in ribbons) + ribbons.discard(ribbon_name) + if ribbons: + raise ValueError('Unknown ribbons: {0}'.format(', '.join(ribbons))) @property def nickname(self): From a2a63938a60a738fae7fd7a7cb1d0dc5bddd1e3f Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 01:00:56 +0200 Subject: [PATCH 11/23] Use helper function for saving --- pokedex/struct/__init__.py | 184 ++++++++++++++---------------- pokedex/struct/_pokemon_struct.py | 2 - 2 files changed, 85 insertions(+), 101 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index 9bf6e5ffc..8ae260e30 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -12,6 +12,7 @@ import base64 import datetime import contextlib +from operator import attrgetter import sqlalchemy.orm.exc @@ -219,19 +220,48 @@ def export_dict(self): """ st = self.structure - def save_trash(result, name, string): + NO_VALUE = object() + def save(target_dict, key, value=NO_VALUE, transform=None, + condition=lambda x: x): + """Set a dict key to a value, if a condition is true + + If value is not given, it is looked up on self. + The value can be transformed by a function before setting. + """ + if value is NO_VALUE: + attrname = key.replace(' ', '_') + value = getattr(self, attrname) + if condition(value): + if transform: + value = transform(value) + target_dict[key] = value + + def save_string(target_dict, string_key, trash_key, string): + """Save a string, including trash bytes""" + target_dict[string_key] = unicode(string) trash = getattr(string, 'original', None) if trash: expected = (string + u'\uffff').encode('utf-16LE') if trash.rstrip('\0') != expected: - result[name] = base64.b64encode(trash) + target_dict[trash_key] = base64.b64encode(trash) + + def save_object(target_dict, key, value=NO_VALUE, **extra): + """Objects are represented as dicts with "name" and a bunch of IDs + + The name is for humans. The ID is the number from the struct. + """ + save(target_dict, key, value=value, transform=lambda value: + dict(name=value.name, **extra)) result = dict( species=dict(id=self.species.id, name=self.species.name), ) - ability = self.ability - if ability: - result['ability'] = dict(id=st.ability_id, name=ability.name) + if self.form != self.species.default_form: + result['form'] = dict(id=st.form_id, name=self.form.form_name) + + save_object(result, 'ability', id=st.ability_id) + save_object(result, 'held item', id=st.held_item_id) + save_object(result, 'pokeball', id=st.dppt_pokeball or st.hgss_pokeball) trainer = dict( id=self.original_trainer_id, @@ -239,86 +269,52 @@ def save_trash(result, name, string): name=unicode(self.original_trainer_name), gender=self.original_trainer_gender ) - save_trash(trainer, 'name trash', self.original_trainer_name) + save_string(trainer, 'name', 'name trash', self.original_trainer_name) if (trainer['id'] or trainer['secret'] or trainer['name'].strip('\0') or trainer['gender'] != 'male'): result['oiginal trainer'] = trainer - if self.form != self.species.default_form: - result['form'] = dict(id=st.form_id, name=self.form.form_name) - if self.held_item: - result['held item'] = dict(id=st.held_item_id, name=self.held_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 != 'special': - result['encounter type'] = self.encounter_type - if self.nickname: - result['nickname'] = unicode(self.nickname) - save_trash(result, 'nickname trash', self.nickname) - if self.egg_location: - result['egg location'] = dict( - id=st.pt_egg_location_id or st.dp_egg_location_id, - name=self.egg_location.name) - elif st.pt_egg_location_id or st.dp_egg_location_id: - result['egg location'] = dict( - id=st.pt_egg_location_id or st.dp_egg_location_id) - if st.dp_egg_location_id: - result['egg location slot'] = 'dp' - if self.met_location: - result['met location'] = dict( - id=st.pt_met_location_id or st.dp_met_location_id, - name=self.met_location.name) - elif st.pt_met_location_id or st.dp_met_location_id: - result['egg location'] = dict( - id=st.pt_met_location_id or st.dp_met_location_id) - if st.dp_met_location_id: - result['met location slot'] = 'dp' - 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=st.dppt_pokeball or st.hgss_pokeball, - name=self.pokeball.name) - if self.met_at_level: - result['met at level'] = self.met_at_level - - if self.is_nicknamed: - result['nicknamed'] = True - if self.is_egg: - result['is egg'] = True - if self.fateful_encounter: - result['fateful encounter'] = True - if self.personality: - result['personality'] = self.personality - if self.gender != 'genderless': - result['gender'] = self.gender - - if st.hidden_ability: - result['has hidden ability'] = st.hidden_ability + save(result, 'exp') + save(result, 'happiness') + save(result, 'markings', transform=sorted) + save(result, 'original country') + save(result, 'original version') + save(result, 'encounter type', condition=lambda et: + (et and et != 'special')) + save_string(result, 'nickname', 'nickname trash', self.nickname) + save(result, 'egg received', self.date_egg_received, + transform=lambda x: x.isoformat()) + save(result, 'date met', + transform=lambda x: x.isoformat()) + save(result, 'pokerus data', self.pokerus) + save(result, 'met at level') + save(result, 'nicknamed', self.is_nicknamed) + save(result, 'is egg') + save(result, 'fateful encounter') + save(result, 'personality') + save(result, 'gender', condition=lambda g: g != 'genderless') + save(result, 'has hidden ability', self.hidden_ability) + save(result, 'ribbons', + sorted(r.replace('_', ' ') for r in self.ribbons)) + + for loc_type in 'egg', 'met': + loc_dict = dict() + save(loc_dict, 'id_pt', st['pt_{0}_location_id'.format(loc_type)]) + save(loc_dict, 'id_dp', st['dp_{0}_location_id'.format(loc_type)]) + save(loc_dict, 'name', + getattr(self, '{0}_location'.format(loc_type)), + transform=attrgetter('name')) + save(result, '{0} location'.format(loc_type), loc_dict) 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 - move['pp'] = st['move%s_pp' % i] - pp_up = st['move%s_pp_ups' % i] - if pp_up: - move['pp ups'] = pp_up - if move: + save(move, 'id', move_object, transform=attrgetter('id')) + save(move, 'name', move_object, transform=attrgetter('name')) + save(move, 'pp ups', st['move%s_pp_ups' % i]) + pp = st['move%s_pp' % i] + if move or pp: + move['pp'] = pp moves.append(move) effort = {} @@ -332,16 +328,10 @@ def save_trash(result, name, string): effort[dct_stat_identifier] = st['effort_' + st_stat_identifier] for contest_stat in 'cool', 'beauty', 'cute', 'smart', 'tough', 'sheen': contest_stats[contest_stat] = st['contest_' + contest_stat] - if any(effort.values()): - result['effort'] = effort - if any(genes.values()): - result['genes'] = genes - if any(contest_stats.values()): - result['contest stats'] = contest_stats - - ribbons = sorted(r.replace('_', ' ') for r in self.ribbons) - if ribbons: - result['ribbons'] = ribbons + save(result, 'effort', effort, condition=any) + save(result, 'genes', genes, condition=any) + save(result, 'contest stats', contest_stats, condition=any) + return result def update(self, dct, **kwargs): @@ -434,18 +424,14 @@ def load_name(attr_name, dct, string_key, trash_key): is_nicknamed='nicknamed', ) for loc_type in 'egg', 'met': - key = '{0} location'.format(loc_type) - if key in dct: + loc_dict = dct.get('{0} location'.format(loc_type)) + if loc_dict: dp_attr = 'dp_{0}_location_id'.format(loc_type) pt_attr = 'pt_{0}_location_id'.format(loc_type) - if dct.get('{0} location slot'.format(loc_type)) == 'dp': - attr = dp_attr - other_attr = pt_attr - else: - attr = pt_attr - other_attr = dp_attr - st[attr] = dct[key]['id'] - st[other_attr] = 0 + if 'id_dp' in loc_dict: + st[dp_attr] = loc_dict['id_dp'] + if 'id_pt' in loc_dict: + st[pt_attr] = loc_dict['id_pt'] delattr(self, '{0}_location'.format(loc_type)) if 'date met' in dct: self.date_met = datetime.datetime.strptime( @@ -798,7 +784,6 @@ def ribbons(self, ribbons): for ribbonset_name in ( 'sinnoh_ribbons', 'hoenn_ribbons', 'sinnoh_contest_ribbons'): ribbonset = self.structure[ribbonset_name] - print ribbonset for ribbon_name in ribbonset: ribbonset[ribbon_name] = (ribbon_name in ribbons) ribbons.discard(ribbon_name) @@ -948,7 +933,8 @@ def update(self, dct, **kwargs): if 'nature' in dct: self.structure.nature_id = dct['nature']['id'] if 'has hidden ability' not in dct: - self.hidden_ability = (self.pokemon.dream_ability == self.ability) + self.hidden_ability = (self.ability == self.pokemon.dream_ability + and self.ability not in self.pokemon.abilities) @cached_property def nature(self): diff --git a/pokedex/struct/_pokemon_struct.py b/pokedex/struct/_pokemon_struct.py index ebb453311..2c7a81e3f 100644 --- a/pokedex/struct/_pokemon_struct.py +++ b/pokedex/struct/_pokemon_struct.py @@ -721,7 +721,6 @@ def make_pokemon_struct(generation): Flag('circle'), ), LeakyEnum(ULInt8('original_country'), - _unset = 0, jp=1, us=2, fr=3, @@ -855,7 +854,6 @@ def make_pokemon_struct(generation): PokemonStringAdapter(String('nickname', 22), 22), Padding(1), LeakyEnum(ULInt8('original_version'), - _unset = 0, sapphire = 1, ruby = 2, emerald = 3, From fcd4c43fab0ae4554179b28b8b913d17c79acd49 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 01:35:56 +0200 Subject: [PATCH 12/23] Prevent SQLA str/unicode warnings by using a Unicode YAML loader --- pokedex/main.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/pokedex/main.py b/pokedex/main.py index c1319dbab..7f760b91d 100644 --- a/pokedex/main.py +++ b/pokedex/main.py @@ -307,6 +307,18 @@ def command_pkm(*args): if options.yaml: import yaml + # Override the default string handling function + # to always return unicode objects. + # Inspired by http://stackoverflow.com/questions/2890146 + # This prevents str/unicode SQLAlchemy warnings. + def construct_yaml_str(self, node): + return self.construct_scalar(node) + class UnicodeLoader(yaml.SafeLoader): + pass + UnicodeLoader.add_constructor(u'tag:yaml.org,2002:str', + construct_yaml_str) + + if not files: # Use sys.stdin in place of name, handle specially later files = [sys.stdin] @@ -319,7 +331,7 @@ def command_pkm(*args): content = f.read() if mode == 'encode': if options.yaml: - dict_ = yaml.load(content) + dict_ = yaml.load(content, Loader=UnicodeLoader) else: dict_ = json.loads(content) struct = cls(session=session, dict_=dict_) From 61b9ffd30e873241b72036f8b7045e23a2522f10 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 01:56:38 +0200 Subject: [PATCH 13/23] Remove InstrumentedList, some docs, restructuring --- pokedex/struct/__init__.py | 111 +++++++++++++++---------------------- 1 file changed, 46 insertions(+), 65 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index 8ae260e30..8e825e661 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -31,11 +31,16 @@ def pokemon_prng(seed): def struct_proxy(name, dependent=[]): + """Proxies to self.structure. + + "blob" is autometically reset by the setter. + The setter deletes all attributes named in ``dependent``. + """ def getter(self): - return getattr(self.structure, name) + return self.structure[name] def setter(self, value): - setattr(self.structure, name, value) + self.structure[name] = value for dep in dependent: delattr(self, dep) del self.blob @@ -44,15 +49,19 @@ def setter(self, value): def struct_frozenset_proxy(name): + """Proxy for sets like ribbons or markings + + "blob" is autometically reset by the setter. + """ def getter(self): - bitstruct = getattr(self.structure, name) + bitstruct = self.structure[name] return frozenset(k for k, v in bitstruct.items() if v) def setter(self, new_set): new_set = set(new_set) - struct = getattr(self.structure, name) + struct = self.structure[name] for key in struct: - setattr(struct, key, key in new_set) + struct[key] = (key in new_set) new_set.discard(key) if new_set: raise ValueError('Unknown values: {0}'.format(', '.join(ribbons))) @@ -62,6 +71,7 @@ def setter(self, new_set): class cached_property(object): + """Caching property. Use del to remove the cache.""" def __init__(self, getter, setter=None): self._getter = getter self._setter = setter @@ -69,14 +79,15 @@ def __init__(self, getter, setter=None): def setter(self, func): """With this setter, the value being set is automatically cached + + "blob" is autometically reset by the setter. """ self._setter = func self.cache_setter_value = True return self def complete_setter(self, func): - """Setter without automatic caching of the set value - """ + """Setter without automatic caching of the set value""" self._setter = func self.cache_setter_value = False return self @@ -114,34 +125,6 @@ def __delete__(self, instance): pass -class InstrumentedList(object): - def __init__(self, callback, initial=()): - self.list = list(initial) - self.callback = callback - - def __getitem__(self, index): - return self.list[index] - - def __setitem__(self, index, value): - self.list[index] = value - self.callback() - - def __delitem__(self, index, value): - self.list[index] = value - self.callback() - - def append(self, item): - self.list.append(item) - self.callback() - - def extend(self, extralist): - self.list.extend(extralist) - self.callback() - - def __iter__(self): - return iter(self.list) - - class SaveFilePokemon(object): u"""Base class for an individual Pokémon, from the game's point of view. @@ -183,7 +166,6 @@ def __init__(self, blob=None, dict_=None, encrypted=False, session=None): else: # Already decrypted self.blob = blob - else: self.blob = '\0' * (32 * 4 + 8) @@ -276,9 +258,14 @@ def save_object(target_dict, key, value=NO_VALUE, **extra): save(result, 'exp') save(result, 'happiness') - save(result, 'markings', transform=sorted) save(result, 'original country') save(result, 'original version') + save(result, 'met at level') + save(result, 'is egg') + save(result, 'fateful encounter') + save(result, 'personality') + + save(result, 'markings', transform=sorted) save(result, 'encounter type', condition=lambda et: (et and et != 'special')) save_string(result, 'nickname', 'nickname trash', self.nickname) @@ -287,11 +274,7 @@ def save_object(target_dict, key, value=NO_VALUE, **extra): save(result, 'date met', transform=lambda x: x.isoformat()) save(result, 'pokerus data', self.pokerus) - save(result, 'met at level') save(result, 'nicknamed', self.is_nicknamed) - save(result, 'is egg') - save(result, 'fateful encounter') - save(result, 'personality') save(result, 'gender', condition=lambda g: g != 'genderless') save(result, 'has hidden ability', self.hidden_ability) save(result, 'ribbons', @@ -328,9 +311,11 @@ def save_object(target_dict, key, value=NO_VALUE, **extra): effort[dct_stat_identifier] = st['effort_' + st_stat_identifier] for contest_stat in 'cool', 'beauty', 'cute', 'smart', 'tough', 'sheen': contest_stats[contest_stat] = st['contest_' + contest_stat] - save(result, 'effort', effort, condition=any) - save(result, 'genes', genes, condition=any) - save(result, 'contest stats', contest_stats, condition=any) + def any_values(d): + return any(d.values()) + save(result, 'effort', effort, condition=any_values) + save(result, 'genes', genes, condition=any_values) + save(result, 'contest stats', contest_stats, condition=any_values) return result @@ -378,7 +363,7 @@ def update(self, dct, **kwargs): if 'pokeball' in dct: self.pokeball = self._get_pokeball(dct['pokeball']['id']) del self.pokeball - def _load_values(source, **values): + def load_values(source, **values): for attrname, key in values.iteritems(): try: value = source[key] @@ -396,14 +381,14 @@ def load_name(attr_name, dct, string_key, trash_key): setattr(self, attr_name, unicode(dct[string_key])) if 'oiginal trainer' in dct: trainer = dct['oiginal trainer'] - _load_values(trainer, + load_values(trainer, original_trainer_id='id', original_trainer_secret_id='secret', original_trainer_gender='gender', ) load_name('original_trainer_name', trainer, 'name', 'name trash') was_nicknamed = self.is_nicknamed - _load_values(dct, + load_values(dct, exp='exp', happiness='happiness', markings='markings', @@ -420,9 +405,7 @@ def load_name(attr_name, dct, string_key, trash_key): ) load_name('nickname', dct, 'nickname', 'nickname trash') self.is_nicknamed = was_nicknamed - _load_values(dct, - is_nicknamed='nicknamed', - ) + load_values(dct, is_nicknamed='nicknamed') for loc_type in 'egg', 'met': loc_dict = dct.get('{0} location'.format(loc_type)) if loc_dict: @@ -465,6 +448,8 @@ def load_name(attr_name, dct, string_key, trash_key): ('contest stats', 'contest')): for name, value in dct.get(key, {}).items(): st['{}_{}'.format(prefix, name.replace(' ', '_'))] = value + del self.stats + del self.blob return self ### Delicious data @@ -539,6 +524,7 @@ def alternate_form(self, alternate_form): else: st.alternate_form_id = 0 del self.form + del self.blob @property def species(self): @@ -711,27 +697,20 @@ def moves(self): .filter(tables.Move.id.in_(move_ids))) moves_dict = dict((move.id, move) for move in move_rows) - def callback(): - def get(index): - try: - return result[x].id - except AttributeError: - return 0 - self.structure.move1_id = get(0) - self.structure.move2_id = get(1) - self.structure.move3_id = get(2) - self.structure.move4_id = get(3) - self._reset() - - result = InstrumentedList( - callback, + result = tuple( [moves_dict.get(move_id, None) for move_id in move_ids]) return result @moves.complete_setter def moves(self, new_moves): - self.moves[:] = new_moves + for i in range(4): + try: + id = new_moves[x].id + except AttributeError: + return 0 + self.structure['move{0}_id'.format(i)] = id + del self.moves @cached_property def move_pp(self): @@ -789,6 +768,7 @@ def ribbons(self, ribbons): ribbons.discard(ribbon_name) if ribbons: raise ValueError('Unknown ribbons: {0}'.format(', '.join(ribbons))) + del self.blob @property def nickname(self): @@ -948,6 +928,7 @@ def nature(self): @nature.setter def nature(self, new_nature): self.structure.nature_id = int(new_nature.game_index) + del self.blob hidden_ability = struct_proxy('hidden_ability') From 5250e889291ccae64a09300fd9bf45aea4ef8a10 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 02:14:47 +0200 Subject: [PATCH 14/23] Fix breakage in gen 4 --- pokedex/struct/__init__.py | 27 +++++++++++++++++++-------- 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index 8e825e661..ed012d3bd 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -235,6 +235,9 @@ def save_object(target_dict, key, value=NO_VALUE, **extra): save(target_dict, key, value=value, transform=lambda value: dict(name=value.name, **extra)) + def any_values(d): + return any(d.values()) + result = dict( species=dict(id=self.species.id, name=self.species.name), ) @@ -243,7 +246,12 @@ def save_object(target_dict, key, value=NO_VALUE, **extra): save_object(result, 'ability', id=st.ability_id) save_object(result, 'held item', id=st.held_item_id) - save_object(result, 'pokeball', id=st.dppt_pokeball or st.hgss_pokeball) + + ball_dict = {} + save(ball_dict, 'id_dppt', st.dppt_pokeball) + save(ball_dict, 'id_hgss', st.hgss_pokeball) + save_object(ball_dict, 'name', self.pokeball) + save(result, 'pokeball', ball_dict, condition=any_values) trainer = dict( id=self.original_trainer_id, @@ -276,7 +284,6 @@ def save_object(target_dict, key, value=NO_VALUE, **extra): save(result, 'pokerus data', self.pokerus) save(result, 'nicknamed', self.is_nicknamed) save(result, 'gender', condition=lambda g: g != 'genderless') - save(result, 'has hidden ability', self.hidden_ability) save(result, 'ribbons', sorted(r.replace('_', ' ') for r in self.ribbons)) @@ -311,8 +318,6 @@ def save_object(target_dict, key, value=NO_VALUE, **extra): effort[dct_stat_identifier] = st['effort_' + st_stat_identifier] for contest_stat in 'cool', 'beauty', 'cute', 'smart', 'tough', 'sheen': contest_stats[contest_stat] = st['contest_' + contest_stat] - def any_values(d): - return any(d.values()) save(result, 'effort', effort, condition=any_values) save(result, 'genes', genes, condition=any_values) save(result, 'contest stats', contest_stats, condition=any_values) @@ -361,7 +366,10 @@ def update(self, dct, **kwargs): st.held_item_id = dct['held item']['id'] del self.held_item if 'pokeball' in dct: - self.pokeball = self._get_pokeball(dct['pokeball']['id']) + if 'id_dppt' in dct['pokeball']: + st.dppt_pokeball = dct['pokeball']['id_dppt'] + if 'id_hgss' in dct['pokeball']: + st.hgss_pokeball = dct['pokeball']['id_hgss'] del self.pokeball def load_values(source, **values): for attrname, key in values.iteritems(): @@ -401,7 +409,6 @@ def load_name(attr_name, dct, string_key, trash_key): fateful_encounter='fateful encounter', gender='gender', personality='personality', - hidden_ability='has hidden ability', ) load_name('nickname', dct, 'nickname', 'nickname trash') self.is_nicknamed = was_nicknamed @@ -595,9 +602,11 @@ def pokeball(self, pokeball): pokeball_id = pokeball.id boundary = 492 - 17 if pokeball_id >= boundary: + st.dppt_pokeball = 0 st.hgss_pokeball = pokeball_id - boundary else: st.dppt_pokeball = pokeball_id + st.hgss_pokeball = 0 @cached_property def egg_location(self): @@ -861,7 +870,7 @@ class SaveFilePokemonGen4(SaveFilePokemon): pokemon_struct = make_pokemon_struct(generation=generation_id) def export_dict(self): - result = super(SaveFilePokemonGen5, self).export_dict() + result = super(SaveFilePokemonGen4, self).export_dict() if any(self.shiny_leaves): result['shiny leaves'] = self.shiny_leaves return result @@ -912,7 +921,9 @@ def update(self, dct, **kwargs): super(SaveFilePokemonGen5, self).update(dct) if 'nature' in dct: self.structure.nature_id = dct['nature']['id'] - if 'has hidden ability' not in dct: + if 'has hidden ability' in dct: + self.hidden_ability = dct['has hidden ability'] + else: self.hidden_ability = (self.ability == self.pokemon.dream_ability and self.ability not in self.pokemon.abilities) From 483ec7cd1a068965a25a5873ff33ef8312674b10 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 13:47:51 +0200 Subject: [PATCH 15/23] Fix up forms and hidden ability flag --- pokedex/struct/__init__.py | 9 ++++++--- pokedex/struct/_pokemon_struct.py | 6 +++--- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index ed012d3bd..1a8bfb5f8 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -242,7 +242,8 @@ def any_values(d): species=dict(id=self.species.id, name=self.species.name), ) if self.form != self.species.default_form: - result['form'] = dict(id=st.form_id, name=self.form.form_name) + result['form'] = dict( + id=st.alternate_form_id, name=self.form.form_name) save_object(result, 'ability', id=st.ability_id) save_object(result, 'held item', id=st.held_item_id) @@ -345,7 +346,7 @@ def update(self, dct, **kwargs): del self.ability reset_form = False if 'form' in dct: - st.alternate_form = dct['form'] + st.alternate_form_id = dct['form']['id'] reset_form = True if 'species' in dct: st.national_id = dct['species']['id'] @@ -434,7 +435,8 @@ def load_name(attr_name, dct, string_key, trash_key): if 'moves' in dct: pp_reset_indices = [] for i, movedict in enumerate(dct['moves']): - st['move{0}_id'.format(i + 1)] = movedict['id'] + if 'id' in movedict: + st['move{0}_id'.format(i + 1)] = movedict['id'] if 'pp' in movedict: st['move{0}_pp'.format(i + 1)] = movedict['pp'] else: @@ -914,6 +916,7 @@ def export_dict(self): if self.nature: result['nature'] = dict( id=self.structure.nature_id, name=self.nature.name) + result['has hidden ability'] = self.hidden_ability return result def update(self, dct, **kwargs): diff --git a/pokedex/struct/_pokemon_struct.py b/pokedex/struct/_pokemon_struct.py index 2c7a81e3f..ece59e299 100644 --- a/pokedex/struct/_pokemon_struct.py +++ b/pokedex/struct/_pokemon_struct.py @@ -19,7 +19,7 @@ pokemon_forms = { # Unown - 201: 'abcdefghijklmnopqrstuvwxyz!?', + 201: list('abcdefghijklmnopqrstuvwxyz') + ['exclamation', 'question'], # Deoxys 386: ['normal', 'attack', 'defense', 'speed'], @@ -33,7 +33,7 @@ 423: ['west', 'east'], # Rotom - 479: ['normal', 'heat', 'wash', 'frost', 'fan', 'cut'], + 479: ['normal', 'heat', 'wash', 'frost', 'fan', 'mow'], # Giratina 487: ['altered', 'origin'], @@ -45,7 +45,7 @@ 493: [ 'normal', 'fighting', 'flying', 'poison', 'ground', 'rock', 'bug', 'ghost', 'steel', 'fire', 'water', 'grass', - 'thunder', 'psychic', 'ice', 'dragon', 'dark', '???', + 'thunder', 'psychic', 'ice', 'dragon', 'dark', 'unknown', ], } From 8483e57c5df804182390523a8a3072a8a661b8bf Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 13:51:40 +0200 Subject: [PATCH 16/23] Don't group ribbons into one attribute --- pokedex/struct/__init__.py | 31 +++++++------------------------ 1 file changed, 7 insertions(+), 24 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index 1a8bfb5f8..0d6c4c32e 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -285,8 +285,9 @@ def any_values(d): save(result, 'pokerus data', self.pokerus) save(result, 'nicknamed', self.is_nicknamed) save(result, 'gender', condition=lambda g: g != 'genderless') - save(result, 'ribbons', - sorted(r.replace('_', ' ') for r in self.ribbons)) + for name in 'sinnoh ribbons', 'sinnoh contest ribbons', 'hoenn ribbons': + save(result, name, transform=lambda ribbons: + sorted(r.replace('_', ' ') for r in ribbons)) for loc_type in 'egg', 'met': loc_dict = dict() @@ -430,8 +431,10 @@ def load_name(attr_name, dct, string_key, trash_key): if 'egg received' in dct: self.date_egg_received = datetime.datetime.strptime( dct['egg received'], '%Y-%m-%d').date() - if 'ribbons' in dct: - self.ribbons = (r.replace(' ', '_') for r in dct['ribbons']) + for name in 'sinnoh ribbons', 'sinnoh contest ribbons', 'hoenn ribbons': + if name in dct: + setattr(self, name.replace(' ', '_'), + (r.replace(' ', '_') for r in dct[name])) if 'moves' in dct: pp_reset_indices = [] for i, movedict in enumerate(dct['moves']): @@ -761,26 +764,6 @@ def move_pp(self, new_pps): 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) - - @ribbons.setter - def ribbons(self, ribbons): - ribbons = set(ribbons) - for ribbonset_name in ( - 'sinnoh_ribbons', 'hoenn_ribbons', 'sinnoh_contest_ribbons'): - ribbonset = self.structure[ribbonset_name] - for ribbon_name in ribbonset: - ribbonset[ribbon_name] = (ribbon_name in ribbons) - ribbons.discard(ribbon_name) - if ribbons: - raise ValueError('Unknown ribbons: {0}'.format(', '.join(ribbons))) - del self.blob - @property def nickname(self): return self.structure.nickname From 840c5f368cae57017cf75f45650fa37692be1634 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 14:49:21 +0200 Subject: [PATCH 17/23] Save & load trash bytes --- pokedex/struct/__init__.py | 11 +++++++++++ pokedex/struct/_pokemon_struct.py | 17 ++++++++--------- 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index 0d6c4c32e..e7a623b61 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -324,6 +324,14 @@ def any_values(d): save(result, 'genes', genes, condition=any_values) save(result, 'contest stats', contest_stats, condition=any_values) + trash = [] + while True: + try: + trash.append(st['trash_{0}'.format(len(trash))]) + except KeyError: + break + save(result, 'trash values', trash, condition=any) + return result def update(self, dct, **kwargs): @@ -460,6 +468,9 @@ def load_name(attr_name, dct, string_key, trash_key): ('contest stats', 'contest')): for name, value in dct.get(key, {}).items(): st['{}_{}'.format(prefix, name.replace(' ', '_'))] = value + if 'trash values' in dct: + for i, data in enumerate(dct['trash values']): + st['trash_{0}'.format(i)] = data del self.stats del self.blob return self diff --git a/pokedex/struct/_pokemon_struct.py b/pokedex/struct/_pokemon_struct.py index ece59e299..70aa0ca80 100644 --- a/pokedex/struct/_pokemon_struct.py +++ b/pokedex/struct/_pokemon_struct.py @@ -687,9 +687,9 @@ def make_pokemon_struct(generation): 5: ULInt8('nature_id'), }[generation] - padding_or_hidden_ability = { - 4: Padding(1), - 5: Flag('hidden_ability'), + hidden_ability_with_padding = { + 4: ULInt16('trash_1'), + 5: Embed(Struct('', Flag('hidden_ability'), ULInt8('trash_1'))), }[generation] PokemonStringAdapter = { @@ -700,7 +700,7 @@ def make_pokemon_struct(generation): return Struct('pokemon_struct', # Header ULInt32('personality'), # XXX aughgh http://bulbapedia.bulbagarden.net/wiki/Personality - Padding(2), + ULInt16('trash_0'), ULInt16('checksum'), # XXX should be checked or calculated # Block A @@ -845,14 +845,13 @@ def make_pokemon_struct(generation): Flag('fateful_encounter'), )), leaves_or_nature, - padding_or_hidden_ability, - Padding(1), + hidden_ability_with_padding, ULInt16('pt_egg_location_id'), ULInt16('pt_met_location_id'), # Block C PokemonStringAdapter(String('nickname', 22), 22), - Padding(1), + ULInt8('trash_2'), LeakyEnum(ULInt8('original_version'), sapphire = 1, ruby = 2, @@ -889,7 +888,7 @@ def make_pokemon_struct(generation): Flag('cool_ribbon_great'), Flag('cool_ribbon'), ), - Padding(4), + ULInt32('trash_3'), # Block D PokemonStringAdapter(String('original_trainer_name', 16), 16), @@ -919,5 +918,5 @@ def make_pokemon_struct(generation): hgss_gift = 24, # starter; fossil; bebe's eevee (pt only??) ), ULInt8('hgss_pokeball'), - Padding(1), + ULInt8('trash_4'), ) From 857aee1a2937937185efd51dde9602e269d6c9b2 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 16:14:41 +0200 Subject: [PATCH 18/23] Add Python format to the pokedex pkm command --- pokedex/main.py | 32 +++++++++++++++++++++++--------- 1 file changed, 23 insertions(+), 9 deletions(-) diff --git a/pokedex/main.py b/pokedex/main.py index 7f760b91d..37e369048 100644 --- a/pokedex/main.py +++ b/pokedex/main.py @@ -5,6 +5,8 @@ import textwrap import json import base64 +import ast +import pprint import pokedex.db import pokedex.db.load @@ -283,9 +285,13 @@ def command_pkm(*args): Options: --gen=NUM, -g Generation to use (4 or 5) --crypt, -c Use encrypted binary format. - --yaml, -y Use YAML as the human-readable format. - Requires the PyYAML library to be installed. - Default is to use JSON. + --format=FORMAT, -f FORMAT + Select the human-readable format to use. + FORMAT can be: + json (default): use JSON. + yaml: use YAML. Needs the PyYAML library + installed. + python: use Python literal syntax --base64, -b Use Base64 encoding for the binary format. By default, on for 'encode' but off for 'decode' --binary, -B Output raw binary data, do not use Base64 @@ -296,7 +302,7 @@ def command_pkm(*args): parser = get_parser(verbose=False) parser.add_option('-g', '--gen', default=5, type=int) parser.add_option('-c', '--crypt', action='store_true') - parser.add_option('-y', '--yaml', action='store_true') + parser.add_option('-f', '--format', default='json') parser.add_option('-b', '--base64', action='store_true', default=(mode == 'encode')) parser.add_option('-B', '--no-base64', action='store_false', dest='base64') @@ -304,7 +310,7 @@ def command_pkm(*args): session = get_session(options) cls = pokedex.struct.save_file_pokemon_classes[options.gen] - if options.yaml: + if options.format == 'yaml': import yaml # Override the default string handling function @@ -330,10 +336,14 @@ class UnicodeLoader(yaml.SafeLoader): with open(filename) as f: content = f.read() if mode == 'encode': - if options.yaml: + if options.format == 'yaml': dict_ = yaml.load(content, Loader=UnicodeLoader) - else: + elif options.format == 'json': dict_ = json.loads(content) + elif options.format == 'python': + dict_ = ast.literal_eval(content) + else: + raise parser.error('Bad "format"') struct = cls(session=session, dict_=dict_) if options.crypt: data = struct.as_encrypted @@ -349,10 +359,14 @@ class UnicodeLoader(yaml.SafeLoader): struct = cls( blob=content, encrypted=options.crypt, session=session) dict_ = struct.export_dict() - if options.yaml: + if options.format == 'yaml': print yaml.safe_dump(dict_, explicit_start=True), - else: + elif options.format == 'json': print json.dumps(dict_), + elif options.format == 'python': + pprint.pprint(dict_) + else: + raise parser.error('Bad "format"') def command_help(): From ba5fa4355a934a696c95e738b51394366007fc65 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 16:15:49 +0200 Subject: [PATCH 19/23] Add tests --- pokedex/tests/test_struct.py | 162 +++++++++++++++++++++++++++++++ scripts/test-struct-roundtrip.py | 125 ++++++++++++++++++++++++ 2 files changed, 287 insertions(+) create mode 100644 pokedex/tests/test_struct.py create mode 100755 scripts/test-struct-roundtrip.py diff --git a/pokedex/tests/test_struct.py b/pokedex/tests/test_struct.py new file mode 100644 index 000000000..bdad70aa4 --- /dev/null +++ b/pokedex/tests/test_struct.py @@ -0,0 +1,162 @@ +# Encoding: utf8 + +import base64 + +import pytest + +from pokedex import struct +from pokedex.db import connect, tables, util + +session = connect() + +def check_with_roundtrip(gen, pkmn, expected): + blob = pkmn.blob + del pkmn.blob + assert blob == pkmn.blob + + assert pkmn.export_dict() == expected + from_dict = struct.save_file_pokemon_classes[5](session=session, + dict_=expected) + assert from_dict.blob == blob + assert from_dict.export_dict() == expected + + from_blob = struct.save_file_pokemon_classes[5](session=session, + blob=pkmn.blob) + assert from_blob.blob == blob + assert from_blob.export_dict() == expected + + +voltorb_species = util.get(session, tables.PokemonSpecies, 'voltorb') +def voltorb_and_dict(): + pkmn = struct.save_file_pokemon_classes[5](session=session) + voltorb_species = util.get(session, tables.PokemonSpecies, 'voltorb') + pkmn.species = voltorb_species + expected = { + 'gender': 'male', + 'species': dict(id=100, name=u'Voltorb'), + 'nickname': u'\0' * 11, + 'nickname trash': 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAA==', + 'moves': [], + } + return pkmn, expected + + +def test_species(): + pkmn, expected = voltorb_and_dict() + assert pkmn.species == voltorb_species + assert pkmn.pokemon == voltorb_species.default_pokemon + assert pkmn.form == voltorb_species.default_form + assert pkmn.export_dict() == expected + + +def test_nickname(): + pkmn, expected = voltorb_and_dict() + pkmn.nickname = pkmn.nickname + expected['nicknamed'] = True + check_with_roundtrip(5, pkmn, expected) + + pkmn.is_nicknamed = False + del expected['nicknamed'] + check_with_roundtrip(5, pkmn, expected) + +def test_moves(): + pkmn, expected = voltorb_and_dict() + new_moves = (util.get(session, tables.Move, 'sonicboom'), ) + expected['moves'] = [dict(id=49, name=u'SonicBoom', pp=0)] + pkmn.moves = new_moves + assert pkmn.moves == new_moves + check_with_roundtrip(5, pkmn, expected) + + new_moves += (util.get(session, tables.Move, 'explosion'),) + expected['moves'].append(dict(id=153, name=u'Explosion', pp=0)) + pkmn.moves = new_moves + assert pkmn.moves == new_moves + check_with_roundtrip(5, pkmn, expected) + + new_pp = (20,) + expected['moves'][0]['pp'] = 20 + pkmn.move_pp = new_pp + assert pkmn.move_pp == (20, 0, 0, 0) + check_with_roundtrip(5, pkmn, expected) + +def test_personality(): + pkmn, expected = voltorb_and_dict() + assert pkmn.is_shiny == True + pkmn.personality = 12345 + assert pkmn.is_shiny == False + expected['personality'] = 12345 + check_with_roundtrip(5, pkmn, expected) + +def test_pokeball(): + pkmn, expected = voltorb_and_dict() + masterball = util.get(session, tables.Item, 'master-ball') + pkmn.pokeball = masterball + assert pkmn.pokeball == masterball + expected['pokeball'] = dict(id_dppt=1, name='Master Ball') + check_with_roundtrip(5, pkmn, expected) + +def test_experience(): + pkmn, expected = voltorb_and_dict() + exp = 2340 + pkmn.exp = exp + assert pkmn.exp == exp + assert pkmn.experience_rung.experience < pkmn.exp + assert pkmn.next_experience_rung.experience > pkmn.exp + assert pkmn.experience_rung.level + 1 == pkmn.next_experience_rung.level + assert (pkmn.experience_rung.growth_rate == + pkmn.next_experience_rung.growth_rate == + pkmn.species.growth_rate) + assert pkmn.level == pkmn.experience_rung.level + assert pkmn.exp_to_next == pkmn.next_experience_rung.experience - pkmn.exp + rung_difference = (pkmn.next_experience_rung.experience - + pkmn.experience_rung.experience) + assert pkmn.progress_to_next == ( + pkmn.exp - pkmn.experience_rung.experience) / float(rung_difference) + expected['exp'] = exp + check_with_roundtrip(5, pkmn, expected) + +def test_ability(): + pkmn, expected = voltorb_and_dict() + ability = util.get(session, tables.Ability, 'drizzle') + pkmn.ability = ability + assert pkmn.ability == ability + expected['ability'] = dict(id=2, name='Drizzle') + check_with_roundtrip(5, pkmn, expected) + +def test_squitle_blob(): + # Japanese Dream World Squirtle, coutresy of + # http://projectpokemon.org/events/ + blob = base64.b64decode('J2ZqBgAAICQHAAAAkOaKyTACAABGLAABAAAAAAAAAAAAAAAAA' + 'AAAACEAJwCRAG4AIx4eKAAAAAD171MHAAAAAAAAAQAAAAAAvDDLMKww4TD//wAAAAAAAA' + 'AAAAD//wAVAAAAAAAAAAAw/zD/T/9S/0f///8AAAAAAAAACgoOAABLAAAZCgAAAA==') + expected = { + 'ability': {'id': 44, 'name': u'Rain Dish'}, + 'date met': '2010-10-14', + 'exp': 560, + 'gender': 'male', + 'genes': {u'attack': 31, + u'defense': 27, + u'hp': 21, + u'special attack': 21, + u'special defense': 3, + u'speed': 7}, + 'happiness': 70, + 'met at level': 10, + 'met location': {'id_dp': 75, 'name': u'Spring Path'}, + 'moves': [{'id': 33, 'name': u'Tackle', 'pp': 35}, + {'id': 39, 'name': u'Tail Whip', 'pp': 30}, + {'id': 145, 'name': u'Bubble', 'pp': 30}, + {'id': 110, 'name': u'Withdraw', 'pp': 40}], + 'nickname': u'ゼニガメ', + 'nickname trash': 'vDDLMKww4TD//wAAAAAAAAAAAAD//w==', + 'oiginal trainer': {'gender': 'male', + 'id': 59024, + 'name': u'ＰＰｏｒｇ', + 'secret': 51594}, + 'original country': 'jp', + 'original version': 21, + 'personality': 107636263, + 'pokeball': {'id_dppt': 25, 'name': u'Hyper Potion'}, + 'species': {'id': 7, 'name': u'Squirtle'}} + pkmn = struct.save_file_pokemon_classes[5](session=session, blob=blob) + check_with_roundtrip(5, pkmn, expected) diff --git a/scripts/test-struct-roundtrip.py b/scripts/test-struct-roundtrip.py new file mode 100755 index 000000000..b8bfec35a --- /dev/null +++ b/scripts/test-struct-roundtrip.py @@ -0,0 +1,125 @@ +#! /usr/bin/env python +""" +This is an ad-hoc testing script. YMMV +""" + +import os +import sys +import pprint +import binascii +import traceback +import subprocess +import tempfile +import itertools + +import yaml # you need to pip install pyyaml +from blessings import Terminal # you need to pip install blessings + +from pokedex import struct +from pokedex.db import connect + +session = connect(engine_args=dict(echo=False)) + +if len(sys.argv) < 1: + print 'Give this script a bunch of PKM files to test roundtrips on.' + print 'A number (e.g. "4") will be interpreted as the generation of' + print 'the following files, until a new generation is given.' + print 'Use "./5" for a file named 5.' + print + print 'If mismatches are found, your screen will be filled with colorful' + print 'reports. You need the colordiff program for this.' + +def printable(c): + if ord(' ') < ord(c) < ord('~'): + return c + else: + return '.' + +def colordiff(str1, str2, prefix='tmp-'): + if str1 != str2: + with tempfile.NamedTemporaryFile(prefix=prefix + '.', suffix='.a') as file1: + with tempfile.NamedTemporaryFile(prefix=prefix + '.', suffix='.b') as file2: + file1.write(str1) + file2.write(str2) + file1.flush() + file2.flush() + p = subprocess.Popen(['colordiff', '-U999', file1.name, file2.name]) + p.communicate() + else: + print prefix, 'match:' + print str1 + +Class = struct.save_file_pokemon_classes[5] + +filenames_left = list(reversed(sys.argv[1:])) + +while filenames_left: + filename = filenames_left.pop() + print filename + + try: + generation = int(filename) + except ValueError: + pass + else: + Class = struct.save_file_pokemon_classes[generation] + continue + + if os.path.isdir(filename): + for name in sorted(os.listdir(filename), reverse=True): + joined = os.path.join(filename, name) + if name.endswith('.pkm') or os.path.isdir(joined): + filenames_left.append(joined) + continue + + with open(filename) as f: + blob = f.read()[:0x88] + + if blob[0] == blob[1] == blob[2] == blob[3] == '\0': + print binascii.hexlify(blob) + print 'Probably not a PKM file' + + try: + orig_object = Class(blob, session=session) + dict_ = orig_object.export_dict() + except Exception: + traceback.print_exc() + print binascii.hexlify(blob) + continue + orig_object.blob + new_object = Class(dict_=dict_, session=session) + try: + blob_again = new_object.blob + dict_again = new_object.export_dict() + except Exception: + colordiff(yaml.dump(orig_object.structure), yaml.dump(new_object.structure), 'struct') + traceback.print_exc() + continue + + if (dict_ != dict_again) or (blob != blob_again): + colordiff(yaml.dump(orig_object.structure), yaml.dump(new_object.structure), 'struct') + colordiff(yaml.safe_dump(dict_), yaml.safe_dump(dict_again), 'yaml') + t = Terminal() + for pass_number in 1, 2, 3: + for i, (a, b) in enumerate(itertools.izip_longest(blob, blob_again, fillvalue='\xbb')): + if (i - 8) % 32 == 0: + # Block boundary + sys.stdout.write(' ') + a_hex = binascii.hexlify(a) + b_hex = binascii.hexlify(b) + if a != b: + if pass_number == 1: + sys.stdout.write(t.green(printable(a))) + sys.stdout.write(t.red(printable(b))) + elif pass_number == 2: + sys.stdout.write(t.green(a_hex)) + elif pass_number == 3: + sys.stdout.write(t.red(b_hex)) + else: + if pass_number == 1: + sys.stdout.write(printable(a)) + sys.stdout.write(printable(b)) + else: + sys.stdout.write(a_hex) + print + print From 271f0455191717a46a90eddd196af8ec37ad0a5f Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 16:16:02 +0200 Subject: [PATCH 20/23] Fix errors found by tests --- pokedex/struct/__init__.py | 51 +++++++++++++++++++++++--------------- 1 file changed, 31 insertions(+), 20 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index e7a623b61..4befb66c4 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -251,7 +251,7 @@ def any_values(d): ball_dict = {} save(ball_dict, 'id_dppt', st.dppt_pokeball) save(ball_dict, 'id_hgss', st.hgss_pokeball) - save_object(ball_dict, 'name', self.pokeball) + save(ball_dict, 'name', self.pokeball, transform=attrgetter('name')) save(result, 'pokeball', ball_dict, condition=any_values) trainer = dict( @@ -445,6 +445,7 @@ def load_name(attr_name, dct, string_key, trash_key): (r.replace(' ', '_') for r in dct[name])) if 'moves' in dct: pp_reset_indices = [] + i = -1 for i, movedict in enumerate(dct['moves']): if 'id' in movedict: st['move{0}_id'.format(i + 1)] = movedict['id'] @@ -591,9 +592,6 @@ def form(self): def form(self, form): self.structure.national_id = form.species.id self.structure.alternate_form = form.form_identifier - del self.species - del self.pokemon - self._reset() @cached_property def pokeball(self): @@ -659,7 +657,7 @@ def level(self): @cached_property def experience_rung(self): growth_rate = self.species.growth_rate - return (session.query(tables.Experience) + return (self.session.query(tables.Experience) .filter(tables.Experience.growth_rate == growth_rate) .filter(tables.Experience.experience <= self.exp) .order_by(tables.Experience.level.desc()) @@ -668,8 +666,9 @@ def experience_rung(self): @cached_property def next_experience_rung(self): level = self.level + growth_rate = self.species.growth_rate if level < 100: - return (session.query(tables.Experience) + return (self.session.query(tables.Experience) .filter(tables.Experience.growth_rate == growth_rate) .filter(tables.Experience.level == level + 1) .one()) @@ -723,18 +722,18 @@ def moves(self): moves_dict = dict((move.id, move) for move in move_rows) result = tuple( - [moves_dict.get(move_id, None) for move_id in move_ids]) + [moves_dict.get(move_id, None) for move_id in move_ids if move_id]) return result @moves.complete_setter def moves(self, new_moves): - for i in range(4): - try: - id = new_moves[x].id - except AttributeError: - return 0 - self.structure['move{0}_id'.format(i)] = id + ( + self.structure.move1_id, + self.structure.move2_id, + self.structure.move3_id, + self.structure.move4_id, + ) = ([m.id for m in new_moves] + [0, 0, 0, 0])[:4] del self.moves @cached_property @@ -748,7 +747,13 @@ def move_pp(self): @move_pp.complete_setter def move_pp(self, new_pps): - self.move_pp[:] = new_pps + ( + self.structure.move1_pp, + self.structure.move2_pp, + self.structure.move3_pp, + self.structure.move4_pp, + ) = (list(new_pps) + [0, 0, 0, 0])[:4] + del self.move_pp original_trainer_id = struct_proxy('original_trainer_id') original_trainer_secret_id = struct_proxy('original_trainer_secret_id') @@ -910,7 +915,10 @@ def export_dict(self): if self.nature: result['nature'] = dict( id=self.structure.nature_id, name=self.nature.name) - result['has hidden ability'] = self.hidden_ability + ability_is_hidden = (self.ability == self.pokemon.dream_ability) + if (ability_is_hidden != bool(self.hidden_ability) or + self.pokemon.dream_ability in self.pokemon.abilities): + result['has hidden ability'] = self.hidden_ability return result def update(self, dct, **kwargs): @@ -918,11 +926,14 @@ def update(self, dct, **kwargs): super(SaveFilePokemonGen5, self).update(dct) if 'nature' in dct: self.structure.nature_id = dct['nature']['id'] - if 'has hidden ability' in dct: - self.hidden_ability = dct['has hidden ability'] - else: - self.hidden_ability = (self.ability == self.pokemon.dream_ability - and self.ability not in self.pokemon.abilities) + if any(x in dct for x in + ('has hidden ability', 'species', 'form', 'ability')): + if 'has hidden ability' in dct: + self.hidden_ability = dct['has hidden ability'] + else: + self.hidden_ability = ( + self.ability == self.pokemon.dream_ability and + self.ability not in self.pokemon.abilities) @cached_property def nature(self): From b944d830b73d0a5c6827e08726c6bcd8cec20762 Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 17:22:23 +0200 Subject: [PATCH 21/23] Test update() too --- pokedex/struct/__init__.py | 26 +++++++---- pokedex/tests/test_struct.py | 91 +++++++++++++++++++++++++----------- 2 files changed, 81 insertions(+), 36 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index 4befb66c4..5876fe5c7 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -283,7 +283,7 @@ def any_values(d): save(result, 'date met', transform=lambda x: x.isoformat()) save(result, 'pokerus data', self.pokerus) - save(result, 'nicknamed', self.is_nicknamed) + result['nicknamed'] = self.is_nicknamed save(result, 'gender', condition=lambda g: g != 'genderless') for name in 'sinnoh ribbons', 'sinnoh contest ribbons', 'hoenn ribbons': save(result, name, transform=lambda ribbons: @@ -334,7 +334,7 @@ def any_values(d): return result - def update(self, dct, **kwargs): + def update(self, dct=None, **kwargs): """Updates the pokemon from a YAML/JSON-compatible dict Dicts that don't specify all the data are allowed. They update the @@ -349,6 +349,8 @@ def update(self, dct, **kwargs): """ st = self.structure session = self.session + if dct is None: + dct = {} dct.update(kwargs) if 'ability' in dct: st.ability_id = dct['ability']['id'] @@ -365,8 +367,7 @@ def update(self, dct, **kwargs): if reset_form: del self.form if not self.is_nicknamed: - self.nickname = self.species.name - self.is_nicknamed = False + del self.nickname if self.species.gender_rate == -1: self.gender = 'genderless' elif self.gender == 'genderless': @@ -405,7 +406,6 @@ def load_name(attr_name, dct, string_key, trash_key): original_trainer_gender='gender', ) load_name('original_trainer_name', trainer, 'name', 'name trash') - was_nicknamed = self.is_nicknamed load_values(dct, exp='exp', happiness='happiness', @@ -421,8 +421,10 @@ def load_name(attr_name, dct, string_key, trash_key): personality='personality', ) load_name('nickname', dct, 'nickname', 'nickname trash') - self.is_nicknamed = was_nicknamed - load_values(dct, is_nicknamed='nicknamed') + if 'nicknamed' in dct: + self.is_nicknamed = dct['nicknamed'] + elif 'nickname' in dct: + self.is_nicknamed = self.nickname != self.species.name for loc_type in 'egg', 'met': loc_dict = dct.get('{0} location'.format(loc_type)) if loc_dict: @@ -791,7 +793,7 @@ def nickname(self, value): del self.blob @nickname.deleter - def nickname(self, value): + def nickname(self): self.structure.nickname = '' self.is_nicknamed = False del self.blob @@ -876,7 +878,9 @@ def export_dict(self): result['shiny leaves'] = self.shiny_leaves return result - def update(self, dct, **kwargs): + def update(self, dct=None, **kwargs): + if dct is None: + dct = {} dct.update(kwargs) if 'shiny leaves' in dct: self.shiny_leaves = dct['shiny leaves'] @@ -921,7 +925,9 @@ def export_dict(self): result['has hidden ability'] = self.hidden_ability return result - def update(self, dct, **kwargs): + def update(self, dct=None, **kwargs): + if dct is None: + dct = {} dct.update(kwargs) super(SaveFilePokemonGen5, self).update(dct) if 'nature' in dct: diff --git a/pokedex/tests/test_struct.py b/pokedex/tests/test_struct.py index bdad70aa4..6107c2cef 100644 --- a/pokedex/tests/test_struct.py +++ b/pokedex/tests/test_struct.py @@ -7,6 +7,8 @@ from pokedex import struct from pokedex.db import connect, tables, util +from pokedex.tests import positional_params + session = connect() def check_with_roundtrip(gen, pkmn, expected): @@ -35,6 +37,7 @@ def voltorb_and_dict(): 'gender': 'male', 'species': dict(id=100, name=u'Voltorb'), 'nickname': u'\0' * 11, + 'nicknamed': False, 'nickname trash': 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAA==', 'moves': [], } @@ -48,57 +51,92 @@ def test_species(): assert pkmn.form == voltorb_species.default_form assert pkmn.export_dict() == expected - -def test_nickname(): - pkmn, expected = voltorb_and_dict() - pkmn.nickname = pkmn.nickname - expected['nicknamed'] = True - check_with_roundtrip(5, pkmn, expected) - - pkmn.is_nicknamed = False - del expected['nicknamed'] - check_with_roundtrip(5, pkmn, expected) - -def test_moves(): +@positional_params([True], [False]) +def test_moves(use_update): pkmn, expected = voltorb_and_dict() new_moves = (util.get(session, tables.Move, 'sonicboom'), ) expected['moves'] = [dict(id=49, name=u'SonicBoom', pp=0)] - pkmn.moves = new_moves + if use_update: + pkmn.update(moves=expected['moves']) + else: + pkmn.moves = new_moves assert pkmn.moves == new_moves check_with_roundtrip(5, pkmn, expected) new_moves += (util.get(session, tables.Move, 'explosion'),) expected['moves'].append(dict(id=153, name=u'Explosion', pp=0)) - pkmn.moves = new_moves + if use_update: + pkmn.update(moves=expected['moves']) + else: + pkmn.moves = new_moves assert pkmn.moves == new_moves check_with_roundtrip(5, pkmn, expected) new_pp = (20,) expected['moves'][0]['pp'] = 20 - pkmn.move_pp = new_pp + if use_update: + pkmn.update(moves=expected['moves']) + else: + pkmn.move_pp = new_pp assert pkmn.move_pp == (20, 0, 0, 0) check_with_roundtrip(5, pkmn, expected) -def test_personality(): +@positional_params([True], [False]) +def test_personality(use_update): pkmn, expected = voltorb_and_dict() assert pkmn.is_shiny == True - pkmn.personality = 12345 + if use_update: + pkmn.update(personality=12345) + else: + pkmn.personality = 12345 assert pkmn.is_shiny == False expected['personality'] = 12345 check_with_roundtrip(5, pkmn, expected) -def test_pokeball(): +@positional_params([True], [False]) +def test_pokeball(use_update): pkmn, expected = voltorb_and_dict() masterball = util.get(session, tables.Item, 'master-ball') - pkmn.pokeball = masterball - assert pkmn.pokeball == masterball expected['pokeball'] = dict(id_dppt=1, name='Master Ball') + if use_update: + pkmn.update(pokeball=expected['pokeball']) + else: + pkmn.pokeball = masterball + assert pkmn.pokeball == masterball + check_with_roundtrip(5, pkmn, expected) + +@positional_params([True], [False]) +def test_nickname(use_update): + pkmn, expected = voltorb_and_dict() + if use_update: + pkmn.update(nickname=unicode(pkmn.nickname)) + else: + pkmn.nickname = pkmn.nickname + expected['nicknamed'] = True + check_with_roundtrip(5, pkmn, expected) + + if use_update: + pkmn.update(nicknamed=False) + else: + pkmn.is_nicknamed = False + expected['nicknamed'] = False + check_with_roundtrip(5, pkmn, expected) + + if use_update: + pkmn.update(nicknamed=True) + else: + pkmn.is_nicknamed = True + expected['nicknamed'] = True check_with_roundtrip(5, pkmn, expected) -def test_experience(): +@positional_params([True], [False]) +def test_experience(use_update): pkmn, expected = voltorb_and_dict() exp = 2340 - pkmn.exp = exp + if use_update: + pkmn.update(exp=exp) + else: + pkmn.exp = exp assert pkmn.exp == exp assert pkmn.experience_rung.experience < pkmn.exp assert pkmn.next_experience_rung.experience > pkmn.exp @@ -115,7 +153,8 @@ def test_experience(): expected['exp'] = exp check_with_roundtrip(5, pkmn, expected) -def test_ability(): +@positional_params([True], [False]) +def test_ability(use_update): pkmn, expected = voltorb_and_dict() ability = util.get(session, tables.Ability, 'drizzle') pkmn.ability = ability @@ -123,9 +162,8 @@ def test_ability(): expected['ability'] = dict(id=2, name='Drizzle') check_with_roundtrip(5, pkmn, expected) -def test_squitle_blob(): - # Japanese Dream World Squirtle, coutresy of - # http://projectpokemon.org/events/ +def test_squirtle_blob(): + # Japanese Dream World Squirtle from http://projectpokemon.org/events blob = base64.b64decode('J2ZqBgAAICQHAAAAkOaKyTACAABGLAABAAAAAAAAAAAAAAAAA' 'AAAACEAJwCRAG4AIx4eKAAAAAD171MHAAAAAAAAAQAAAAAAvDDLMKww4TD//wAAAAAAAA' 'AAAAD//wAVAAAAAAAAAAAw/zD/T/9S/0f///8AAAAAAAAACgoOAABLAAAZCgAAAA==') @@ -149,6 +187,7 @@ def test_squitle_blob(): {'id': 110, 'name': u'Withdraw', 'pp': 40}], 'nickname': u'ゼニガメ', 'nickname trash': 'vDDLMKww4TD//wAAAAAAAAAAAAD//w==', + 'nicknamed': False, 'oiginal trainer': {'gender': 'male', 'id': 59024, 'name': u'ＰＰｏｒｇ', From 3f513653e12f086914216567f9ad564776c2f07b Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 17:47:09 +0200 Subject: [PATCH 22/23] Make level settable, export level And if exp is just enough to reach the given level, leave exp out --- pokedex/struct/__init__.py | 19 +++++++++++- pokedex/tests/test_struct.py | 60 +++++++++++++++++++++++++----------- 2 files changed, 60 insertions(+), 19 deletions(-) diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index 5876fe5c7..89469c0eb 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -265,7 +265,6 @@ def any_values(d): trainer['name'].strip('\0') or trainer['gender'] != 'male'): result['oiginal trainer'] = trainer - save(result, 'exp') save(result, 'happiness') save(result, 'original country') save(result, 'original version') @@ -274,6 +273,10 @@ def any_values(d): save(result, 'fateful encounter') save(result, 'personality') + save(result, 'level') + if self.exp != self.experience_rung.experience: + save(result, 'exp') + save(result, 'markings', transform=sorted) save(result, 'encounter type', condition=lambda et: (et and et != 'special')) @@ -420,6 +423,12 @@ def load_name(attr_name, dct, string_key, trash_key): gender='gender', personality='personality', ) + if 'level' in dct: + if 'exp' in dct: + if self.level != dct['level']: + raise ValueError('level and exp not compatible') + else: + self.level = dct['level'] load_name('nickname', dct, 'nickname', 'nickname trash') if 'nicknamed' in dct: self.is_nicknamed = dct['nicknamed'] @@ -656,6 +665,14 @@ def met_location(self): def level(self): return self.experience_rung.level + @level.setter + def level(self, level): + growth_rate = self.species.growth_rate + self.exp = (self.session.query(tables.Experience) + .filter(tables.Experience.growth_rate == growth_rate) + .filter(tables.Experience.level == level) + .one().experience) + @cached_property def experience_rung(self): growth_rate = self.species.growth_rate diff --git a/pokedex/tests/test_struct.py b/pokedex/tests/test_struct.py index 6107c2cef..ef37fec41 100644 --- a/pokedex/tests/test_struct.py +++ b/pokedex/tests/test_struct.py @@ -36,6 +36,7 @@ def voltorb_and_dict(): expected = { 'gender': 'male', 'species': dict(id=100, name=u'Voltorb'), + 'level': 1, 'nickname': u'\0' * 11, 'nicknamed': False, 'nickname trash': 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAA==', @@ -132,25 +133,48 @@ def test_nickname(use_update): @positional_params([True], [False]) def test_experience(use_update): pkmn, expected = voltorb_and_dict() - exp = 2340 + for exp in 2197, 2200: + if use_update: + pkmn.update(exp=exp) + else: + pkmn.exp = exp + assert pkmn.exp == exp + assert pkmn.experience_rung.experience <= pkmn.exp + assert pkmn.next_experience_rung.experience > pkmn.exp + assert pkmn.experience_rung.level + 1 == pkmn.next_experience_rung.level + assert (pkmn.experience_rung.growth_rate == + pkmn.next_experience_rung.growth_rate == + pkmn.species.growth_rate) + assert pkmn.level == pkmn.experience_rung.level + assert pkmn.exp_to_next == pkmn.next_experience_rung.experience - pkmn.exp + rung_difference = (pkmn.next_experience_rung.experience - + pkmn.experience_rung.experience) + assert pkmn.progress_to_next == ( + pkmn.exp - pkmn.experience_rung.experience) / float(rung_difference) + if exp == 2197: + expected['level'] = 13 + else: + expected['exp'] = exp + expected['level'] = 13 + check_with_roundtrip(5, pkmn, expected) + +def test_update_inconsistent_exp_level(): + pkmn, expected = voltorb_and_dict() + with pytest.raises(ValueError): + pkmn.update(exp=0, level=100) + +@positional_params([True], [False]) +def test_level(use_update): + pkmn, expected = voltorb_and_dict() + level = 10 if use_update: - pkmn.update(exp=exp) + pkmn.update(level=level) else: - pkmn.exp = exp - assert pkmn.exp == exp - assert pkmn.experience_rung.experience < pkmn.exp - assert pkmn.next_experience_rung.experience > pkmn.exp - assert pkmn.experience_rung.level + 1 == pkmn.next_experience_rung.level - assert (pkmn.experience_rung.growth_rate == - pkmn.next_experience_rung.growth_rate == - pkmn.species.growth_rate) - assert pkmn.level == pkmn.experience_rung.level - assert pkmn.exp_to_next == pkmn.next_experience_rung.experience - pkmn.exp - rung_difference = (pkmn.next_experience_rung.experience - - pkmn.experience_rung.experience) - assert pkmn.progress_to_next == ( - pkmn.exp - pkmn.experience_rung.experience) / float(rung_difference) - expected['exp'] = exp + pkmn.level = level + assert pkmn.level == level + assert pkmn.experience_rung.level == level + assert pkmn.experience_rung.experience == pkmn.exp + expected['level'] = level check_with_roundtrip(5, pkmn, expected) @positional_params([True], [False]) @@ -170,7 +194,6 @@ def test_squirtle_blob(): expected = { 'ability': {'id': 44, 'name': u'Rain Dish'}, 'date met': '2010-10-14', - 'exp': 560, 'gender': 'male', 'genes': {u'attack': 31, u'defense': 27, @@ -179,6 +202,7 @@ def test_squirtle_blob(): u'special defense': 3, u'speed': 7}, 'happiness': 70, + 'level': 10, 'met at level': 10, 'met location': {'id_dp': 75, 'name': u'Spring Path'}, 'moves': [{'id': 33, 'name': u'Tackle', 'pp': 35}, From e33d25e3df6b88f43fb2ddd4fc9422eea2622a1e Mon Sep 17 00:00:00 2001 From: Petr Viktorin Date: Sun, 14 Oct 2012 18:17:22 +0200 Subject: [PATCH 23/23] pokede pkm: Make the base64 encoding default more sane --- pokedex/main.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/pokedex/main.py b/pokedex/main.py index 37e369048..954ef1347 100644 --- a/pokedex/main.py +++ b/pokedex/main.py @@ -284,7 +284,6 @@ def command_pkm(*args): Options: --gen=NUM, -g Generation to use (4 or 5) - --crypt, -c Use encrypted binary format. --format=FORMAT, -f FORMAT Select the human-readable format to use. FORMAT can be: @@ -292,9 +291,11 @@ def command_pkm(*args): yaml: use YAML. Needs the PyYAML library installed. python: use Python literal syntax + --crypt, -c Use encrypted binary format. --base64, -b Use Base64 encoding for the binary format. - By default, on for 'encode' but off for 'decode' - --binary, -B Output raw binary data, do not use Base64 + --binary, -B Output raw binary data. This is the default, + but you need to specify -B explicitly if you're + dumping binary data to a terminal. If no files are given, reads from standard input. """).encode(sys.getdefaultencoding(), 'replace') @@ -303,8 +304,7 @@ def command_pkm(*args): parser.add_option('-g', '--gen', default=5, type=int) parser.add_option('-c', '--crypt', action='store_true') parser.add_option('-f', '--format', default='json') - parser.add_option('-b', '--base64', action='store_true', - default=(mode == 'encode')) + parser.add_option('-b', '--base64', action='store_true', default=None) parser.add_option('-B', '--no-base64', action='store_false', dest='base64') options, files = parser.parse_args(list(args[1:])) @@ -324,6 +324,18 @@ class UnicodeLoader(yaml.SafeLoader): UnicodeLoader.add_constructor(u'tag:yaml.org,2002:str', construct_yaml_str) + if options.format not in ('yaml', 'json', 'python'): + raise parser.error('Bad "format"') + + if mode == 'encode' and options.base64 is None: + try: + isatty = sys.stdout.isatty + except AttributeError: + pass + else: + if isatty(): + parser.error('Refusing to dump binary data to terminal. ' + 'Please use -B to override, or -b for base64.') if not files: # Use sys.stdin in place of name, handle specially later @@ -342,8 +354,6 @@ class UnicodeLoader(yaml.SafeLoader): dict_ = json.loads(content) elif options.format == 'python': dict_ = ast.literal_eval(content) - else: - raise parser.error('Bad "format"') struct = cls(session=session, dict_=dict_) if options.crypt: data = struct.as_encrypted @@ -365,8 +375,6 @@ class UnicodeLoader(yaml.SafeLoader): print json.dumps(dict_), elif options.format == 'python': pprint.pprint(dict_) - else: - raise parser.error('Bad "format"') def command_help():