Add the pkm subcommand to pokedex

pokedex/main.py

@@ -2,11 +2,15 @@
 from optparse import OptionParser
 import os
 import sys
+import textwrap
+import json
+import base64
 
 import pokedex.db
 import pokedex.db.load
 import pokedex.db.tables
 import pokedex.lookup
+import pokedex.struct
 from pokedex import defaults
 
 def main():
@@ -260,6 +264,85 @@ def command_lookup(*args):
             print
 
 
+def command_pkm(*args):
+    if args and args[0] == 'encode':
+        mode = 'encode'
+    elif args and args[0] == 'decode':
+        mode = 'decode'
+    else:
+        print textwrap.dedent(u"""
+            Convert binary Pokémon data (aka PKM files) to/from JSON/YAML.
+            usage: pokedex pkm (encode|decode) [options] <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)
+                --crypt, -c    Use encrypted binary format.
+                --yaml, -y     Use YAML as the human-readable format.
+                               Requires the PyYAML library to be installed.
+                               Default is to use JSON.
+                --base64, -b   Use Base64 encoding for the binary format.
+                               By default, on for 'encode' but off for 'decode'
+                --binary, -B   Output raw binary data, do not use Base64
+
+            If no files are given, reads from standard input.
+            """).encode(sys.getdefaultencoding(), 'replace')
+        return
+    parser = get_parser(verbose=False)
+    parser.add_option('-g', '--gen', default=5, type=int)
+    parser.add_option('-c', '--crypt', action='store_true')
+    parser.add_option('-y', '--yaml', action='store_true')
+    parser.add_option('-b', '--base64', action='store_true',
+        default=(mode == 'encode'))
+    parser.add_option('-B', '--no-base64', action='store_false', dest='base64')
+    options, files = parser.parse_args(list(args[1:]))
+
+    session = get_session(options)
+    cls = pokedex.struct.save_file_pokemon_classes[options.gen]
+    if options.yaml:
+        import yaml
+
+    if not files:
+        # Use sys.stdin in place of name, handle specially later
+        files = [sys.stdin]
+
+    for filename in files:
+        if filename is sys.stdin:
+            content = sys.stdin.read()
+        else:
+            with open(filename) as f:
+                content = f.read()
+        if mode == 'encode':
+            if options.yaml:
+                dict_ = yaml.load(content)
+            else:
+                dict_ = json.loads(content)
+            struct = cls(session=session, dict_=dict_)
+            if options.crypt:
+                data = struct.as_encrypted
+            else:
+                data = struct.as_struct
+            if options.base64:
+                print base64.b64encode(data)
+            else:
+                sys.stdout.write(data)
+        else:
+            if options.base64:
+                content = base64.b64decode(content)
+            struct = cls(
+                blob=content, encrypted=options.crypt, session=session)
+            dict_ = struct.export_dict()
+            if options.yaml:
+                print yaml.safe_dump(dict_, explicit_start=True),
+            else:
+                print json.dumps(dict_),
+
+
 def command_help():
     print u"""pokedex -- a command-line Pokédex interface
 usage: pokedex {command} [options...]
@@ -269,6 +352,7 @@ See http://bugs.veekun.com/projects/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.