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): diff --git a/pokedex/main.py b/pokedex/main.py index fe4557730..954ef1347 100644 --- a/pokedex/main.py +++ b/pokedex/main.py @@ -2,11 +2,17 @@ from optparse import OptionParser import os import sys +import textwrap +import json +import base64 +import ast +import pprint 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 +266,117 @@ 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) + --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 + --crypt, -c Use encrypted binary format. + --base64, -b Use Base64 encoding for the binary format. + --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') + 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('-f', '--format', default='json') + 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:])) + + session = get_session(options) + cls = pokedex.struct.save_file_pokemon_classes[options.gen] + if options.format == '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 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 + 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.format == 'yaml': + dict_ = yaml.load(content, Loader=UnicodeLoader) + elif options.format == 'json': + dict_ = json.loads(content) + elif options.format == 'python': + dict_ = ast.literal_eval(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.format == 'yaml': + print yaml.safe_dump(dict_, explicit_start=True), + elif options.format == 'json': + print json.dumps(dict_), + elif options.format == 'python': + pprint.pprint(dict_) + + def command_help(): print u"""pokedex -- a command-line Pokédex interface usage: pokedex {command} [options...] @@ -269,6 +386,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. diff --git a/pokedex/struct/__init__.py b/pokedex/struct/__init__.py index b6a3e736a..89469c0eb 100644 --- a/pokedex/struct/__init__.py +++ b/pokedex/struct/__init__.py @@ -9,11 +9,18 @@ """ import struct +import base64 +import datetime +import contextlib +from operator import attrgetter -from pokedex.db import tables +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 pokemon_struct +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.""" @@ -23,42 +30,152 @@ def pokemon_prng(seed): yield seed >> 16 +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 self.structure[name] + + def setter(self, value): + self.structure[name] = value + for dep in dependent: + delattr(self, dep) + del self.blob + + return property(getter, setter) + + +def struct_frozenset_proxy(name): + """Proxy for sets like ribbons or markings + + "blob" is autometically reset by the setter. + """ + def getter(self): + 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 = self.structure[name] + for key in struct: + 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) + + +class cached_property(object): + """Caching property. Use del to remove the cache.""" + 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 + + "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""" + 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 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, 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 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 + else: + self.blob = '\0' * (32 * 4 + 8) + if session: + self.session = session else: - # Already decrypted - self.blob = blob + self.session = None - self.structure = pokemon_struct.parse(self.blob) + if dict_: + self.update(dict_) @property def as_struct(self): @@ -80,6 +197,296 @@ def as_encrypted(self): # Stuff back into a string, and done return struct.pack(struct_def, *shuffled) + def export_dict(self): + """Exports the pokemon as a YAML/JSON-compatible dict + """ + st = self.structure + + 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: + 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)) + + def any_values(d): + return any(d.values()) + + result = dict( + species=dict(id=self.species.id, name=self.species.name), + ) + if self.form != self.species.default_form: + 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) + + ball_dict = {} + save(ball_dict, 'id_dppt', st.dppt_pokeball) + save(ball_dict, 'id_hgss', st.hgss_pokeball) + save(ball_dict, 'name', self.pokeball, transform=attrgetter('name')) + save(result, 'pokeball', ball_dict, condition=any_values) + + trainer = dict( + id=self.original_trainer_id, + secret=self.original_trainer_secret_id, + name=unicode(self.original_trainer_name), + gender=self.original_trainer_gender + ) + 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 + + save(result, 'happiness') + 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, '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')) + 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) + 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: + sorted(r.replace('_', ' ') for r in 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 = {} + 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 = {} + genes = {} + contest_stats = {} + 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('-', ' ') + 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': + contest_stats[contest_stat] = st['contest_' + contest_stat] + save(result, 'effort', effort, condition=any_values) + 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=None, **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 + if dct is None: + dct = {} + dct.update(kwargs) + if 'ability' in dct: + st.ability_id = dct['ability']['id'] + del self.ability + reset_form = False + if 'form' in dct: + st.alternate_form_id = dct['form']['id'] + 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: + del self.nickname + 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.held_item_id = dct['held item']['id'] + del self.held_item + if 'pokeball' in dct: + 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(): + try: + value = source[key] + except KeyError: + 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: + trainer = dct['oiginal 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') + load_values(dct, + exp='exp', + happiness='happiness', + markings='markings', + original_country='original country', + original_version='original version', + encounter_type='encounter type', + pokerus='pokerus data', + met_at_level='met at level', + is_egg='is egg', + fateful_encounter='fateful encounter', + 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'] + 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: + dp_attr = 'dp_{0}_location_id'.format(loc_type) + pt_attr = 'pt_{0}_location_id'.format(loc_type) + 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( + dct['date met'], '%Y-%m-%d').date() + if 'egg received' in dct: + self.date_egg_received = datetime.datetime.strptime( + dct['egg received'], '%Y-%m-%d').date() + 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 = [] + i = -1 + for i, movedict in enumerate(dct['moves']): + 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: + 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 + 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 + 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 + 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 + ### Delicious data @property def is_shiny(self): @@ -98,46 +505,22 @@ 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 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 - st = self.structure - self._pokemon = session.query(tables.Pokemon).get(st.national_id) - self._pokemon_form = session.query(tables.PokemonForm) \ - .with_parent(self._pokemon) \ - .filter_by(name=st.alternate_form) \ - .one() - self._ability = self._session.query(tables.Ability).get(st.ability_id) - - growth_rate = self._pokemon.evolution_chain.growth_rate - self._experience_rung = session.query(tables.Experience) \ - .filter(tables.Experience.growth_rate == growth_rate) \ - .filter(tables.Experience.experience <= st.exp) \ - .order_by(tables.Experience.level.desc()) \ - [0] - level = self._experience_rung.level - - self._next_experience_rung = None - if level < 100: - self._next_experience_rung = session.query(tables.Experience) \ - .filter(tables.Experience.growth_rate == growth_rate) \ - .filter(tables.Experience.level == level + 1) \ - .one() - - self._held_item = None - if st.held_item_id: - self._held_item = session.query(tables.ItemGameIndex) \ - .filter_by(generation_id = 4, game_index = st.held_item_id).one().item - - self._stats = [] - for pokemon_stat in self._pokemon.stats: - structure_name = pokemon_stat.stat.name.lower().replace(' ', '_') - gene = st.ivs['iv_' + structure_name] - exp = st['effort_' + structure_name] - - if pokemon_stat.stat.name == u'HP': + @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 @@ -155,106 +538,224 @@ def use_database_session(self, session): ), ) - self._stats.append(stat_tup) - + stats.append(stat_tup) + return tuple(stats) - 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 alternate_form(self): + st = self.structure + forms = pokemon_forms.get(st.national_id) + if forms: + return forms[st.alternate_form_id] + else: + return None - self._moves = [moves_dict.get(move_id, None) for move_id in move_ids] + @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: + st.alternate_form_id = 0 + del self.form + del self.blob - if st.hgss_pokeball >= 17: - pokeball_id = st.hgss_pokeball - 17 + 492 + @property + def species(self): + if self.form: + return self.form.species else: - pokeball_id = st.dppt_pokeball - self._pokeball = session.query(tables.ItemGameIndex) \ - .filter_by(generation_id = 4, game_index = pokeball_id).one().item + return None - 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 + @species.setter + def species(self, species): + self.form = species.default_form - 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 + @property + def pokemon(self): + if self.form: + return self.form.pokemon + else: + return None - self._met_location = session.query(tables.LocationGameIndex) \ - .filter_by(generation_id = 4, game_index = met_loc_id).one().location + @pokemon.setter + def pokemon(self, pokemon): + self.form = pokemon.default_form - @property - def species(self): - # XXX forme! - return self._pokemon + @cached_property + def form(self): + st = self.structure + session = self.session + if st.national_id: + pokemon = session.query(tables.Pokemon).get(st.national_id) + 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 - @property - def species_form(self): - return self._pokemon_form + @form.setter + def form(self, form): + self.structure.national_id = form.species.id + self.structure.alternate_form = form.form_identifier - @property + @cached_property def pokeball(self): - return self._pokeball + st = self.structure + if st.hgss_pokeball >= 17: + pokeball_id = st.hgss_pokeball - 17 + 492 + elif st.dppt_pokeball: + pokeball_id = st.dppt_pokeball + else: + return None + return self._get_pokeball(pokeball_id) - @property + 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.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): - return self._egg_location + st = self.structure + egg_loc_id = st.pt_egg_location_id or st.dp_egg_location_id + if egg_loc_id: + 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 - @property + @cached_property def met_location(self): - return self._met_location - - @property - def shiny_leaves(self): - return ( - self.structure.shining_leaves.leaf1, - self.structure.shining_leaves.leaf2, - self.structure.shining_leaves.leaf3, - self.structure.shining_leaves.leaf4, - self.structure.shining_leaves.leaf5, - ) + st = self.structure + met_loc_id = st.pt_met_location_id or st.dp_met_location_id + if met_loc_id: + 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 @property def level(self): - return self._experience_rung.level + 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 + 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()) + [0]) + + @cached_property + def next_experience_rung(self): + level = self.level + growth_rate = self.species.growth_rate + if level < 100: + return (self.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 self.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 + result = tuple( + [moves_dict.get(move_id, None) for move_id in move_ids if move_id]) - @property - def moves(self): - return self._moves + return result - @property + @moves.complete_setter + def moves(self, new_moves): + ( + 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 def move_pp(self): return ( self.structure.move1_pp, @@ -263,6 +764,56 @@ def move_pp(self): self.structure.move4_pp, ) + @move_pp.complete_setter + def move_pp(self, 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') + 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') + personality = struct_proxy('personality') + + 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 nickname(self): + return self.structure.nickname + + @nickname.setter + def nickname(self, value): + self.structure.nickname = value + self.is_nicknamed = True + del self.blob + + @nickname.deleter + def nickname(self): + self.structure.nickname = '' + self.is_nicknamed = False + del self.blob ### Utility methods @@ -319,3 +870,112 @@ def reciprocal_crypt(cls, words): words[i] ^= next(prng) return + + @cached_property + def blob(self): + 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 + 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_dict(self): + result = super(SaveFilePokemonGen4, self).export_dict() + if any(self.shiny_leaves): + result['shiny leaves'] = self.shiny_leaves + return result + + 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'] + super(SaveFilePokemonGen4, self).update(dct) + + @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, + 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 + del self.blob + + +class SaveFilePokemonGen5(SaveFilePokemon): + generation_id = 5 + pokemon_struct = make_pokemon_struct(generation=generation_id) + + 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) + 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=None, **kwargs): + if dct is None: + dct = {} + dct.update(kwargs) + super(SaveFilePokemonGen5, self).update(dct) + if 'nature' in dct: + self.structure.nature_id = dct['nature']['id'] + 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): + 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) + del self.blob + + hidden_ability = struct_proxy('hidden_ability') + + +save_file_pokemon_classes = { + 4: SaveFilePokemonGen4, + 5: SaveFilePokemonGen5, +} diff --git a/pokedex/struct/_pokemon_struct.py b/pokedex/struct/_pokemon_struct.py index 79ff0e0f4..70aa0ca80 100644 --- a/pokedex/struct/_pokemon_struct.py +++ b/pokedex/struct/_pokemon_struct.py @@ -17,8 +17,40 @@ # - higher-level validation; see XXXes below # - personality indirectly influences IVs due to PRNG use +pokemon_forms = { + # Unown + 201: list('abcdefghijklmnopqrstuvwxyz') + ['exclamation', 'question'], + + # 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', 'mow'], + + # 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', 'unknown', + ], +} + # The entire gen 4 character table: -character_table = { +character_table_gen4 = { 0x0002: u'ぁ', 0x0003: u'あ', 0x0004: u'ぃ', @@ -465,10 +497,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): @@ -487,25 +573,64 @@ def LittleEndianBitStruct(*args): resizer=lambda _: _, ) + +class StringWithOriginal(unicode): + pass + + 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. + + 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(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): - #padded_text = (obj + u'\xffff' + '\x00' * 12) - padded_text = obj - decoded_text = padded_text.translate(inverse_character_table) - return decoded_text.encode('utf16') + try: + original = obj.original + except AttributeError: + 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): + 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. @@ -527,289 +652,271 @@ 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'], +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) - # 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 _encode(self, obj, context): + return self.values.get(obj, obj) def _decode(self, obj, context): - try: - forms = self.pokemon_forms[ context['national_id'] ] - except KeyError: - return None - - return forms[obj >> 3] + return self.inverted_values.get(obj, obj) - def _encode(self, obj, context): - try: - forms = self.pokemon_forms[ context['national_id'] ] - except KeyError: - return None - 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 +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] + + hidden_ability_with_padding = { + 4: ULInt16('trash_1'), + 5: Embed(Struct('', Flag('hidden_ability'), ULInt8('trash_1'))), + }[generation] + + PokemonStringAdapter = { + 4: PokemonStringAdapterGen4, + 5: PokemonStringAdapterGen5, + }[generation] + + return Struct('pokemon_struct', + # Header + ULInt32('personality'), # XXX aughgh http://bulbapedia.bulbagarden.net/wiki/Personality + ULInt16('trash_0'), + 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'), + ), + LeakyEnum(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', + 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'), + ), -# 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', - 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', - 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 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'), + + Embed(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'), + ), + Embed(EmbeddedBitStruct( + BitField('alternate_form_id', 5), + Enum(BitField('gender', 2), + genderless = 2, + male = 0, + female = 1, + ), + Flag('fateful_encounter'), + )), + leaves_or_nature, + hidden_ability_with_padding, + ULInt16('pt_egg_location_id'), + ULInt16('pt_met_location_id'), + + # Block C + PokemonStringAdapter(String('nickname', 22), 22), + ULInt8('trash_2'), + LeakyEnum(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'), ), - 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, + ULInt32('trash_3'), + + # Block D + PokemonStringAdapter(String('original_trainer_name', 16), 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), ), - 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), -) + LeakyEnum(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'), + ULInt8('trash_4'), + ) diff --git a/pokedex/tests/test_struct.py b/pokedex/tests/test_struct.py new file mode 100644 index 000000000..ef37fec41 --- /dev/null +++ b/pokedex/tests/test_struct.py @@ -0,0 +1,225 @@ +# Encoding: utf8 + +import base64 + +import pytest + +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): + 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'), + 'level': 1, + 'nickname': u'\0' * 11, + 'nicknamed': False, + '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 + +@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)] + 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)) + 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 + 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) + +@positional_params([True], [False]) +def test_personality(use_update): + pkmn, expected = voltorb_and_dict() + assert pkmn.is_shiny == True + 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) + +@positional_params([True], [False]) +def test_pokeball(use_update): + pkmn, expected = voltorb_and_dict() + masterball = util.get(session, tables.Item, 'master-ball') + 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) + +@positional_params([True], [False]) +def test_experience(use_update): + pkmn, expected = voltorb_and_dict() + 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(level=level) + else: + 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]) +def test_ability(use_update): + 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_squirtle_blob(): + # Japanese Dream World Squirtle from 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', + '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, + 'level': 10, + '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==', + 'nicknamed': False, + '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