diff --git a/pokedex/data/csv/natures.csv b/pokedex/data/csv/natures.csv
index d5864af..81a27fe 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 85d45c8..a4c6426 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 fd29593..c42909d 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(*argv):
@@ -266,6 +272,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] <file> ...
+
+ 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...]
@@ -275,6 +392,7 @@ See https://github.com/veekun/pokedex/wiki/CLI for more documentation.
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 b6a3e73..89469c0 100644
--- a/pokedex/struct/__init__.py
+++ b/pokedex/struct/__init__.py
@@ -9,11 +9,18 @@ derived.
"""
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.<name>
+
+ "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")
- # 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)
+ 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)
+
+ else:
+ # Already decrypted
+ self.blob = blob
else:
- # Already decrypted
- self.blob = blob
+ self.blob = '\0' * (32 * 4 + 8)
- self.structure = pokemon_struct.parse(self.blob)
+ if session:
+ self.session = session
+ else:
+ self.session = None
+
+ if dict_:
+ self.update(dict_)
@property
def as_struct(self):
@@ -80,6 +197,296 @@ class SaveFilePokemon(object):
# 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 @@ class SaveFilePokemon(object):
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)
+ @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]
- 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':
+ if pokemon_stat.stat.identifier == u'hp':
calc = calculated_hp
else:
calc = calculated_stat
@@ -155,106 +538,224 @@ class SaveFilePokemon(object):
),
)
- self._stats.append(stat_tup)
+ 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
+ @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
+
+ @property
+ def species(self):
+ if self.form:
+ return self.form.species
+ else:
+ return None
+
+ @species.setter
+ def species(self, species):
+ self.form = species.default_form
+
+ @property
+ def pokemon(self):
+ if self.form:
+ return self.form.pokemon
+ else:
+ return None
+
+ @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)
+ 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
+
+ @form.setter
+ def form(self, form):
+ self.structure.national_id = form.species.id
+ self.structure.alternate_form = form.form_identifier
+
+ @cached_property
+ def pokeball(self):
+ 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)
+
+ 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):
+ 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
+
+ @cached_property
+ 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:
+ 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
+
+ @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.exp
+ else:
+ return 0
+
+ @property
+ def progress_to_next(self):
+ 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
+
+ @cached_property
+ def ability(self):
+ return self.session.query(tables.Ability).get(self.structure.ability_id)
+
+ @ability.setter
+ def ability(self, ability):
+ self.structure.ability_id = ability.id
+
+ @cached_property
+ def held_item(self):
+ 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))
+ move_rows = (self.session.query(tables.Move)
+ .filter(tables.Move.id.in_(move_ids)))
moves_dict = dict((move.id, move) for move in move_rows)
- self._moves = [moves_dict.get(move_id, None) for move_id in move_ids]
+ result = tuple(
+ [moves_dict.get(move_id, None) for move_id in move_ids if move_id])
- if st.hgss_pokeball >= 17:
- pokeball_id = st.hgss_pokeball - 17 + 492
- else:
- pokeball_id = st.dppt_pokeball
- self._pokeball = session.query(tables.ItemGameIndex) \
- .filter_by(generation_id = 4, game_index = pokeball_id).one().item
+ return result
- 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
+ @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
- self._egg_location = None
- if egg_loc_id:
- self._egg_location = session.query(tables.LocationGameIndex) \
- .filter_by(generation_id = 4, game_index = egg_loc_id).one().location
-
- self._met_location = session.query(tables.LocationGameIndex) \
- .filter_by(generation_id = 4, game_index = met_loc_id).one().location
-
- @property
- def species(self):
- # XXX forme!
- return self._pokemon
-
- @property
- def species_form(self):
- return self._pokemon_form
-
- @property
- def pokeball(self):
- return self._pokeball
-
- @property
- def egg_location(self):
- return self._egg_location
-
- @property
- def met_location(self):
- return self._met_location
-
- @property
- def shiny_leaves(self):
- return (
- self.structure.shining_leaves.leaf1,
- self.structure.shining_leaves.leaf2,
- self.structure.shining_leaves.leaf3,
- self.structure.shining_leaves.leaf4,
- self.structure.shining_leaves.leaf5,
- )
-
- @property
- def level(self):
- return self._experience_rung.level
-
- @property
- def exp_to_next(self):
- if self._next_experience_rung:
- return self._next_experience_rung.experience - self.structure.exp
- else:
- return 0
-
- @property
- def progress_to_next(self):
- if self._next_experience_rung:
- return 1.0 \
- * (self.structure.exp - self._experience_rung.experience) \
- / (self._next_experience_rung.experience - self._experience_rung.experience)
- else:
- return 0.0
-
- @property
- def ability(self):
- return self._ability
-
- @property
- def held_item(self):
- return self._held_item
-
- @property
- def stats(self):
- return self._stats
-
- @property
- def moves(self):
- return self._moves
-
- @property
+ @cached_property
def move_pp(self):
return (
self.structure.move1_pp,
@@ -263,6 +764,56 @@ class SaveFilePokemon(object):
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 @@ class SaveFilePokemon(object):
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 79ff0e0..70aa0ca 100644
--- a/pokedex/struct/_pokemon_struct.py
+++ b/pokedex/struct/_pokemon_struct.py
@@ -17,8 +17,40 @@ from construct import *
# - 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 @@ character_table = {
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 @@ class DateAdapter(Adapter):
y, m, d = obj.year - 2000, obj.month, obj.day
return ''.join(chr(n) for n in (y, m, d))
-class PokemonFormAdapter(Adapter):
- """Converts form ids to form names, and vice versa."""
- pokemon_forms = {
- # Unown
- 201: 'abcdefghijklmnopqrstuvwxyz!?',
-
- # Deoxys
- 386: ['normal', 'attack', 'defense', 'speed'],
-
- # Burmy and Wormadam
- 412: ['plant', 'sandy', 'trash'],
- 413: ['plant', 'sandy', 'trash'],
-
- # Shellos and Gastrodon
- 422: ['west', 'east'],
- 423: ['west', 'east'],
-
- # Rotom
- 479: ['normal', 'heat', 'wash', 'frost', 'fan', 'cut'],
-
- # Giratina
- 487: ['altered', 'origin'],
-
- # Shaymin
- 492: ['land', 'sky'],
-
- # Arceus
- 493: [
- 'normal', 'fighting', 'flying', 'poison', 'ground', 'rock',
- 'bug', 'ghost', 'steel', 'fire', 'water', 'grass',
- 'thunder', 'psychic', 'ice', 'dragon', 'dark', '???',
- ],
- }
-
- def _decode(self, obj, context):
- try:
- forms = self.pokemon_forms[ context['national_id'] ]
- except KeyError:
- return None
-
- return forms[obj >> 3]
+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):
- try:
- forms = self.pokemon_forms[ context['national_id'] ]
- except KeyError:
- return None
+ return self.values.get(obj, obj)
- return forms.index(obj) << 3
+ def _decode(self, obj, context):
+ return self.inverted_values.get(obj, obj)
-
-# 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
+# http://projectpokemon.org/wiki/Pokemon_Black/White_NDS_Structure
+# http://projectpokemon.org/forums/showthread.php?11474-Hex-Values-and-Trashbytes-in-B-W#post93598
- # 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,
- ),
+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]
- # 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'),
+ hidden_ability_with_padding = {
+ 4: ULInt16('trash_1'),
+ 5: Embed(Struct('', Flag('hidden_ability'), ULInt8('trash_1'))),
+ }[generation]
- ULInt8('contest_cool'),
- ULInt8('contest_beauty'),
- ULInt8('contest_cute'),
- ULInt8('contest_smart'),
- ULInt8('contest_tough'),
- ULInt8('contest_sheen'),
+ PokemonStringAdapter = {
+ 4: PokemonStringAdapterGen4,
+ 5: PokemonStringAdapterGen5,
+ }[generation]
- 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'),
- ),
+ 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 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 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'),
),
- 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,
+ LeakyEnum(ULInt8('original_country'),
+ jp=1,
+ us=2,
+ fr=3,
+ it=4,
+ de=5,
+ es=7,
+ kr=8,
),
- 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),
-)
+
+ # 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'),
+
+ 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'),
+ ),
+ 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),
+ ),
+ 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 0000000..ef37fec
--- /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 0000000..b8bfec3
--- /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