Merge e33d25e3df into 8cba4ac996

pokedex/data/csv/natures.csv

@@ -1,26 +1,26 @@
-id,identifier,decreased_stat_id,increased_stat_id,hates_flavor_id,likes_flavor_id
+id,identifier,decreased_stat_id,increased_stat_id,hates_flavor_id,likes_flavor_id,game_index
-1,hardy,2,2,1,1
+1,hardy,2,2,1,1,0
-2,bold,2,3,1,5
+2,bold,2,3,1,5,5
-3,modest,2,4,1,2
+3,modest,2,4,1,2,15
-4,calm,2,5,1,4
+4,calm,2,5,1,4,20
-5,timid,2,6,1,3
+5,timid,2,6,1,3,10
-6,lonely,3,2,5,1
+6,lonely,3,2,5,1,1
-7,docile,3,3,5,5
+7,docile,3,3,5,5,6
-8,mild,3,4,5,2
+8,mild,3,4,5,2,16
-9,gentle,3,5,5,4
+9,gentle,3,5,5,4,21
-10,hasty,3,6,5,3
+10,hasty,3,6,5,3,11
-11,adamant,4,2,2,1
+11,adamant,4,2,2,1,3
-12,impish,4,3,2,5
+12,impish,4,3,2,5,8
-13,bashful,4,4,2,2
+13,bashful,4,4,2,2,18
-14,careful,4,5,2,4
+14,careful,4,5,2,4,23
-15,rash,5,4,4,2
+15,rash,5,4,4,2,19
-16,jolly,4,6,2,3
+16,jolly,4,6,2,3,13
-17,naughty,5,2,4,1
+17,naughty,5,2,4,1,4
-18,lax,5,3,4,5
+18,lax,5,3,4,5,9
-19,quirky,5,5,4,4
+19,quirky,5,5,4,4,24
-20,naive,5,6,4,3
+20,naive,5,6,4,3,14
-21,brave,6,2,3,1
+21,brave,6,2,3,1,2
-22,relaxed,6,3,3,5
+22,relaxed,6,3,3,5,7
-23,quiet,6,4,3,2
+23,quiet,6,4,3,2,17
-24,sassy,6,5,3,4
+24,sassy,6,5,3,4,22
-25,serious,6,6,3,3
+25,serious,6,6,3,3,12

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):

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.

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()
+# Generation 5 uses UCS-16, with a few exceptions
-inverse_character_table = dict()
+character_table_gen5 = {
-for in_, out in character_table.iteritems():
+    # Here nintendo just didn't do their homework:
-    inverse_character_table[ord(out)] = in_
+    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
+    u"""Base adapter for names
-    String struct.
+
+    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)
+        try:
-        padded_text = obj
+            original = obj.original
-        decoded_text = padded_text.translate(inverse_character_table)
+        except AttributeError:
-        return decoded_text.encode('utf16')
+            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):
+class LeakyEnum(Adapter):
-    """Converts form ids to form names, and vice versa."""
+    """An Enum that allows unknown values"""
-    pokemon_forms = {
+    def __init__(self, sub, **values):
-        # Unown
+        super(LeakyEnum, self).__init__(sub)
-        201: 'abcdefghijklmnopqrstuvwxyz!?',
+        self.values = values
-
+        self.inverted_values = dict((v, k) for k, v in values.items())
-        # Deoxys
+        assert len(values) == len(self.inverted_values)
-        386: ['normal', 'attack', 'defense', 'speed'],
-
-        # Burmy and Wormadam
-        412: ['plant', 'sandy', 'trash'],
-        413: ['plant', 'sandy', 'trash'],
-
-        # Shellos and Gastrodon
-        422: ['west', 'east'],
-        423: ['west', 'east'],
-
-        # Rotom
-        479: ['normal', 'heat', 'wash', 'frost', 'fan', 'cut'],
-
-        # Giratina
-        487: ['altered', 'origin'],
-
-        # Shaymin
-        492: ['land', 'sky'],
-
-        # Arceus
-        493: [
-            'normal', 'fighting', 'flying', 'poison', 'ground', 'rock',
-            'bug', 'ghost', 'steel', 'fire', 'water', 'grass',
-            'thunder', 'psychic', 'ice', 'dragon', 'dark', '???',
-        ],
-    }
-
-    def _decode(self, obj, context):
-        try:
-            forms = self.pokemon_forms[ context['national_id'] ]
-        except KeyError:
-            return None
-
-        return forms[obj >> 3]
 
     def _encode(self, obj, context):
-        try:
+        return self.values.get(obj, obj)
-            forms = self.pokemon_forms[ context['national_id'] ]
-        except KeyError:
-            return None
 
-        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',
+# http://projectpokemon.org/wiki/Pokemon_Black/White_NDS_Structure
-    # Header
+# http://projectpokemon.org/forums/showthread.php?11474-Hex-Values-and-Trashbytes-in-B-W#post93598
-    ULInt32('personality'),  # XXX aughgh http://bulbapedia.bulbagarden.net/wiki/Personality
-    Padding(2),
-    ULInt16('checksum'),  # XXX should be checked or calculated
 
-    # Block A
+def make_pokemon_struct(generation):
-    ULInt16('national_id'),
+    """Make a pokemon struct class for the given generation
-    ULInt16('held_item_id'),
+    """
-    ULInt16('original_trainer_id'),
+    leaves_or_nature = {
-    ULInt16('original_trainer_secret_id'),
+        4: BitStruct('shining_leaves',
-    ULInt32('exp'),
+                Padding(2),
-    ULInt8('happiness'),
+                Flag('crown'),
-    ULInt8('ability_id'),  # XXX needs to match personality + species
+                Flag('leaf5'),
-    BitStruct('markings',
+                Flag('leaf4'),
-        Padding(2),
+                Flag('leaf3'),
-        Flag('diamond'),
+                Flag('leaf2'),
-        Flag('star'),
+                Flag('leaf1'),
-        Flag('heart'),
+            ),
-        Flag('square'),
+        5: ULInt8('nature_id'),
-        Flag('triangle'),
+    }[generation]
-        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
+    hidden_ability_with_padding = {
-    ULInt8('effort_hp'),
+        4: ULInt16('trash_1'),
-    ULInt8('effort_attack'),
+        5: Embed(Struct('', Flag('hidden_ability'), ULInt8('trash_1'))),
-    ULInt8('effort_defense'),
+    }[generation]
-    ULInt8('effort_speed'),
-    ULInt8('effort_special_attack'),
-    ULInt8('effort_special_defense'),
 
-    ULInt8('contest_cool'),
+    PokemonStringAdapter = {
-    ULInt8('contest_beauty'),
+        4: PokemonStringAdapterGen4,
-    ULInt8('contest_cute'),
+        5: PokemonStringAdapterGen5,
-    ULInt8('contest_smart'),
+    }[generation]
-    ULInt8('contest_tough'),
-    ULInt8('contest_sheen'),
 
-    LittleEndianBitStruct('sinnoh_ribbons',
+    return Struct('pokemon_struct',
-        Padding(4),
+        # Header
-        Flag('premier_ribbon'),
+        ULInt32('personality'),  # XXX aughgh http://bulbapedia.bulbagarden.net/wiki/Personality
-        Flag('classic_ribbon'),
+        ULInt16('trash_0'),
-        Flag('carnival_ribbon'),
+        ULInt16('checksum'),  # XXX should be checked or calculated
-        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
+        # Block A
-    ULInt16('move1_id'),
+        ULInt16('national_id'),
-    ULInt16('move2_id'),
+        ULInt16('held_item_id'),
-    ULInt16('move3_id'),
+        ULInt16('original_trainer_id'),
-    ULInt16('move4_id'),
+        ULInt16('original_trainer_secret_id'),
-    ULInt8('move1_pp'),
+        ULInt32('exp'),
-    ULInt8('move2_pp'),
+        ULInt8('happiness'),
-    ULInt8('move3_pp'),
+        ULInt8('ability_id'),  # XXX needs to match personality + species
-    ULInt8('move4_pp'),
+        BitStruct('markings',
-    ULInt8('move1_pp_ups'),
+            Padding(2),
-    ULInt8('move2_pp_ups'),
+            Flag('diamond'),
-    ULInt8('move3_pp_ups'),
+            Flag('star'),
-    ULInt8('move4_pp_ups'),
+            Flag('heart'),
-
+            Flag('square'),
-    LittleEndianBitStruct('ivs',
+            Flag('triangle'),
-        Flag('is_nicknamed'),
+            Flag('circle'),
-        Flag('is_egg'),
-        BitField('iv_special_defense', 5),
-        BitField('iv_special_attack', 5),
-        BitField('iv_speed', 5),
-        BitField('iv_defense', 5),
-        BitField('iv_attack', 5),
-        BitField('iv_hp', 5),
-    ),
-    LittleEndianBitStruct('hoenn_ribbons',
-        Flag('world_ribbon'),
-        Flag('earth_ribbon'),
-        Flag('national_ribbon'),
-        Flag('country_ribbon'),
-        Flag('sky_ribbon'),
-        Flag('land_ribbon'),
-        Flag('marine_ribbon'),
-        Flag('effort_ribbon'),
-        Flag('artist_ribbon'),
-        Flag('victory_ribbon'),
-        Flag('winning_ribbon'),
-        Flag('champion_ribbon'),
-        Flag('tough_ribbon_master'),
-        Flag('tough_ribbon_hyper'),
-        Flag('tough_ribbon_super'),
-        Flag('tough_ribbon'),
-        Flag('smart_ribbon_master'),
-        Flag('smart_ribbon_hyper'),
-        Flag('smart_ribbon_super'),
-        Flag('smart_ribbon'),
-        Flag('cute_ribbon_master'),
-        Flag('cute_ribbon_hyper'),
-        Flag('cute_ribbon_super'),
-        Flag('cute_ribbon'),
-        Flag('beauty_ribbon_master'),
-        Flag('beauty_ribbon_hyper'),
-        Flag('beauty_ribbon_super'),
-        Flag('beauty_ribbon'),
-        Flag('cool_ribbon_master'),
-        Flag('cool_ribbon_hyper'),
-        Flag('cool_ribbon_super'),
-        Flag('cool_ribbon'),
-    ),
-    EmbeddedBitStruct(
-        PokemonFormAdapter(BitField('alternate_form', 5)),
-        Enum(BitField('gender', 2),
-            genderless = 2,
-            male = 0,
-            female = 1,
         ),
-        Flag('fateful_encounter'),
+        LeakyEnum(ULInt8('original_country'),
-    ),
+            jp=1,
-    BitStruct('shining_leaves',
+            us=2,
-        Padding(2),
+            fr=3,
-        Flag('crown'),
+            it=4,
-        Flag('leaf5'),
+            de=5,
-        Flag('leaf4'),
+            es=7,
-        Flag('leaf3'),
+            kr=8,
-        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,
         ),
-        BitField('met_at_level', 7),
+
-    ),
+        # XXX sum cannot surpass 510
-    Enum(ULInt8('encounter_type'),
+        ULInt8('effort_hp'),
-        special = 0,        # egg; pal park; event; honey tree; shaymin
+        ULInt8('effort_attack'),
-        grass = 2,          # or darkrai
+        ULInt8('effort_defense'),
-        dialga_palkia = 4,
+        ULInt8('effort_speed'),
-        cave = 5,           # or giratina or hall of origin
+        ULInt8('effort_special_attack'),
-        water = 7,
+        ULInt8('effort_special_defense'),
-        building = 9,
+
-        safari_zone = 10,   # includes great marsh
+        ULInt8('contest_cool'),
-        gift = 12,          # starter; fossil; ingame trade?
+        ULInt8('contest_beauty'),
-        # distortion_world = ???,
+        ULInt8('contest_cute'),
-        hgss_gift = 24,     # starter; fossil; bebe's eevee  (pt only??)
+        ULInt8('contest_smart'),
-    ),
+        ULInt8('contest_tough'),
-    ULInt8('hgss_pokeball'),
+        ULInt8('contest_sheen'),
-    Padding(1),
+
-)
+        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'),
+    )

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)

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