Port the CLI to argparse

pokedex/main.py

@@ -1,8 +1,7 @@
 # encoding: utf8
 from __future__ import print_function
 
-import locale
-from optparse import OptionParser
+import argparse
 import os
 import sys
 
@@ -12,47 +11,111 @@ import pokedex.db.tables
 import pokedex.lookup
 from pokedex import defaults
 
-def main(*argv):
-    if len(argv) <= 1:
+
+def main(junk, *argv):
+    if len(argv) <= 0:
         command_help()
+        return
 
-    command = argv[1]
-    args = argv[2:]
+    parser = create_parser()
+    args = parser.parse_args(argv)
+    args.func(parser, args)
 
-    # XXX there must be a better way to get Unicode argv
-    # XXX this doesn't work on Windows durp
-    enc = sys.stdin.encoding or 'utf8'
-    args = [_.decode(enc) if isinstance(_, bytes) else _ for _ in args]
-
-    # Find the command as a function in this file
-    func = globals().get("command_%s" % command, None)
-    if func:
-        func(*args)
-    else:
-        command_help()
 
 def setuptools_entry():
     main(*sys.argv)
 
 
-def get_parser(verbose=True):
-    """Returns an OptionParser prepopulated with the global options.
-
-    `verbose` is whether or not the options should be verbose by default.
+def create_parser():
+    """Build and return an ArgumentParser.
     """
-    parser = OptionParser()
-    parser.add_option('-e', '--engine', dest='engine_uri', default=None)
-    parser.add_option('-i', '--index', dest='index_dir', default=None)
-    parser.add_option('-q', '--quiet', dest='verbose', default=verbose, action='store_false')
-    parser.add_option('-v', '--verbose', dest='verbose', default=verbose, action='store_true')
+    parser = argparse.ArgumentParser(prog='pokedex', description=u'A command-line Pokédex interface')
+    parser.add_argument(
+        '-e', '--engine', dest='engine_uri', default=None,
+        help=u'By default, all commands try to use a SQLite database '
+            u'in the pokedex install directory.  Use this option (or '
+            u'a POKEDEX_DB_ENGINE environment variable) to specify an '
+            u'alternate database.',
+        )
+    parser.add_argument(
+        '-i', '--index', dest='index_dir', default=None,
+        help=u'By default, all commands try to put the lookup index in '
+            u'the pokedex install directory.  Use this option (or a '
+            u'POKEDEX_INDEX_DIR environment variable) to specify an '
+            u'alternate loction.',
+    )
+    parser.add_argument(
+        '-q', '--quiet', dest='verbose', action='store_false',
+        help=u'Don\'t print system output.  This is the default for '
+            'non-system commands and setup.',
+    )
+    parser.add_argument(
+        '-v', '--verbose', dest='verbose', default=False, action='store_true',
+        help=u'Print system output.  This is the default for system '
+            u'commands, except setup.',
+    )
+
+    cmds = parser.add_subparsers(title='Commands')
+    cmd_help = cmds.add_parser('help', help=u'Display this message')
+    cmd_help.set_defaults(func=command_help)
+
+    cmd_lookup = cmds.add_parser('lookup', help=u'Look up something in the Pokédex')
+    cmd_lookup.set_defaults(func=command_lookup)
+    cmd_lookup.add_argument('criteria', nargs='+')
+
+    cmd_load = cmds.add_parser('load', help=u'Load Pokédex data into a database from CSV files')
+    cmd_load.set_defaults(func=command_load, verbose=True)
+    # TODO get the actual default here
+    cmd_load.add_argument(
+        '-d', '--directory', dest='directory', default=None,
+        help="directory containing the CSV files to load")
+    cmd_load.add_argument(
+        '-D', '--drop-tables', dest='drop_tables', default=False, action='store_true',
+        help="drop all tables before loading data")
+    cmd_load.add_argument(
+        '-r', '--recursive', dest='recursive', default=False, action='store_true',
+        help="load and drop all dependent tables (default is to use exactly the given list)")
+    cmd_load.add_argument(
+        '-S', '--safe', dest='safe', default=False, action='store_true',
+        help="disable database-specific optimizations, such as Postgres's COPY FROM")
+    # TODO need a custom handler for splittin' all of these
+    cmd_load.add_argument(
+        '-l', '--langs', dest='langs', default=None,
+        help="comma-separated list of language codes to load, or 'none' (default: all)")
+    cmd_load.add_argument(
+        'tables', nargs='*',
+        help="list of database tables to load (default: all)")
+
+    cmd_dump = cmds.add_parser('dump', help=u'Dump Pokédex data from a database into CSV files')
+    cmd_dump.set_defaults(func=command_dump, verbose=True)
+    cmd_dump.add_argument(
+        '-d', '--directory', dest='directory', default=None,
+        help="directory to place the dumped CSV files")
+    cmd_dump.add_argument(
+        '-l', '--langs', dest='langs', default=None,
+        help="comma-separated list of language codes to load, 'none', or 'all' (default: en)")
+    cmd_dump.add_argument(
+        'tables', nargs='*',
+        help="list of database tables to load (default: all)")
+
+    cmd_reindex = cmds.add_parser('reindex', help=u'Rebuild the lookup index from the database')
+    cmd_reindex.set_defaults(func=command_reindex, verbose=True)
+
+    cmd_setup = cmds.add_parser('setup', help=u'Combine load and reindex')
+    cmd_setup.set_defaults(func=command_setup, verbose=False)
+
+    cmd_status = cmds.add_parser('status', help=u'Print which engine, index, and csv directory would be used for other commands')
+    cmd_status.set_defaults(func=command_status, verbose=True)
+
     return parser
 
-def get_session(options):
+
+def get_session(args):
     """Given a parsed options object, connects to the database and returns a
     session.
     """
 
-    engine_uri = options.engine_uri
+    engine_uri = args.engine_uri
     got_from = 'command line'
 
     if engine_uri is None:
@@ -60,13 +123,14 @@ def get_session(options):
 
     session = pokedex.db.connect(engine_uri)
 
-    if options.verbose:
+    if args.verbose:
         print("Connected to database %(engine)s (from %(got_from)s)"
             % dict(engine=session.bind.url, got_from=got_from))
 
     return session
 
-def get_lookup(options, session=None, recreate=False):
+
+def get_lookup(args, session=None, recreate=False):
     """Given a parsed options object, opens the whoosh index and returns a
     PokedexLookup object.
     """
@@ -74,13 +138,13 @@ def get_lookup(options, session=None, recreate=False):
     if recreate and not session:
         raise ValueError("get_lookup() needs an explicit session to regen the index")
 
-    index_dir = options.index_dir
+    index_dir = args.index_dir
     got_from = 'command line'
 
     if index_dir is None:
         index_dir, got_from = defaults.get_default_index_dir_with_origin()
 
-    if options.verbose:
+    if args.verbose:
         print("Opened lookup index %(index_dir)s (from %(got_from)s)"
             % dict(index_dir=index_dir, got_from=got_from))
 
@@ -91,13 +155,14 @@ def get_lookup(options, session=None, recreate=False):
 
     return lookup
 
-def get_csv_directory(options):
+
+def get_csv_directory(args):
     """Prints and returns the csv directory we're about to use."""
 
-    if not options.verbose:
+    if not args.verbose:
         return
 
-    csvdir = options.directory
+    csvdir = args.directory
     got_from = 'command line'
 
     if csvdir is None:
@@ -111,99 +176,78 @@ def get_csv_directory(options):
 
 ### Plumbing commands
 
-def command_dump(*args):
-    parser = get_parser(verbose=True)
-    parser.add_option('-d', '--directory', dest='directory', default=None)
-    parser.add_option('-l', '--langs', dest='langs', default=None,
-        help="Comma-separated list of languages to dump all strings for. "
-            "Default is English ('en')")
-    options, tables = parser.parse_args(list(args))
+def command_dump(parser, args):
+    session = get_session(args)
+    get_csv_directory(args)
 
-    session = get_session(options)
-    get_csv_directory(options)
-
-    if options.langs is not None:
-        langs = [l.strip() for l in options.langs.split(',')]
+    if args.langs is not None:
+        langs = [l.strip() for l in args.langs.split(',')]
     else:
         langs = None
 
-    pokedex.db.load.dump(session, directory=options.directory,
-                                  tables=tables,
-                                  verbose=options.verbose,
-                                  langs=langs)
+    pokedex.db.load.dump(
+        session,
+        directory=args.directory,
+        tables=args.tables,
+        verbose=args.verbose,
+        langs=langs,
+    )
 
-def command_load(*args):
-    parser = get_parser(verbose=True)
-    parser.add_option('-d', '--directory', dest='directory', default=None)
-    parser.add_option('-D', '--drop-tables', dest='drop_tables', default=False, action='store_true')
-    parser.add_option('-r', '--recursive', dest='recursive', default=False, action='store_true')
-    parser.add_option('-S', '--safe', dest='safe', default=False, action='store_true',
-        help="Do not use backend-specific optimalizations.")
-    parser.add_option('-l', '--langs', dest='langs', default=None,
-        help="Comma-separated list of extra languages to load, or 'none' for none. "
-            "Default is to load 'em all. Example: 'fr,de'")
-    options, tables = parser.parse_args(list(args))
 
-    if not options.engine_uri:
+def command_load(parser, args):
+    if not args.engine_uri:
         print("WARNING: You're reloading the default database, but not the lookup index.  They")
         print("         might get out of sync, and pokedex commands may not work correctly!")
         print("To fix this, run `pokedex reindex` when this command finishes.  Or, just use")
         print("`pokedex setup` to do both at once.")
         print()
 
-    if options.langs == 'none':
+    if args.langs == 'none':
         langs = []
-    elif options.langs is None:
+    elif args.langs is None:
         langs = None
     else:
-        langs = [l.strip() for l in options.langs.split(',')]
+        langs = [l.strip() for l in args.langs.split(',')]
 
-    session = get_session(options)
-    get_csv_directory(options)
+    session = get_session(args)
+    get_csv_directory(args)
 
-    pokedex.db.load.load(session, directory=options.directory,
-                                  drop_tables=options.drop_tables,
-                                  tables=tables,
-                                  verbose=options.verbose,
-                                  safe=options.safe,
-                                  recursive=options.recursive,
-                                  langs=langs)
+    pokedex.db.load.load(
+        session,
+        directory=args.directory,
+        drop_tables=args.drop_tables,
+        tables=args.tables,
+        verbose=args.verbose,
+        safe=args.safe,
+        recursive=args.recursive,
+        langs=langs,
+    )
 
-def command_reindex(*args):
-    parser = get_parser(verbose=True)
-    options, _ = parser.parse_args(list(args))
-
-    session = get_session(options)
-    lookup = get_lookup(options, session=session, recreate=True)
 
+def command_reindex(parser, args):
+    session = get_session(args)
+    get_lookup(args, session=session, recreate=True)
     print("Recreated lookup index.")
 
 
-def command_setup(*args):
-    parser = get_parser(verbose=False)
-    options, _ = parser.parse_args(list(args))
+def command_setup(parser, args):
+    args.directory = None
 
-    options.directory = None
-
-    session = get_session(options)
-    get_csv_directory(options)
-    pokedex.db.load.load(session, directory=None, drop_tables=True,
-                                  verbose=options.verbose,
-                                  safe=False)
-
-    lookup = get_lookup(options, session=session, recreate=True)
+    session = get_session(args)
+    get_csv_directory(args)
+    pokedex.db.load.load(
+        session, directory=None, drop_tables=True,
+        verbose=args.verbose, safe=False)
 
+    get_lookup(args, session=session, recreate=True)
     print("Recreated lookup index.")
 
 
-def command_status(*args):
-    parser = get_parser(verbose=True)
-    options, _ = parser.parse_args(list(args))
-    options.verbose = True
-    options.directory = None
+def command_status(parser, args):
+    args.directory = None
 
     # Database, and a lame check for whether it's been inited at least once
-    session = get_session(options)
+    session = get_session(args)
     print("  - OK!  Connected successfully.")
 
     if pokedex.db.tables.Pokemon.__table__.exists(session.bind):
@@ -212,7 +256,7 @@ def command_status(*args):
         print("  - WARNING: Database appears to be empty.")
 
     # CSV; simple checks that the dir exists
-    csvdir = get_csv_directory(options)
+    csvdir = get_csv_directory(args)
     if not os.path.exists(csvdir):
         print("  - ERROR: No such directory!")
     elif not os.path.isdir(csvdir):
@@ -233,20 +277,17 @@ def command_status(*args):
 
     # Index; the PokedexLookup constructor covers most tests and will
     # cheerfully bomb if they fail
-    lookup = get_lookup(options, recreate=False)
+    get_lookup(args, recreate=False)
     print("  - OK!  Opened successfully.")
 
 
 ### User-facing commands
 
-def command_lookup(*args):
-    parser = get_parser(verbose=False)
-    options, words = parser.parse_args(list(args))
+def command_lookup(parser, args):
+    name = u' '.join(args.criteria)
 
-    name = u' '.join(words)
-
-    session = get_session(options)
-    lookup = get_lookup(options, session=session, recreate=False)
+    session = get_session(args)
+    lookup = get_lookup(args, session=session, recreate=False)
 
     results = lookup.lookup(name)
     if not results:
@@ -269,62 +310,8 @@ def command_lookup(*args):
             print()
 
 
-def command_help():
-    print(u"""pokedex -- a command-line Pokédex interface
-usage: pokedex {command} [options...]
-Run `pokedex setup` first, or nothing will work!
-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.
-
-System commands:
-    load                Load Pokédex data into a database from CSV files.
-    dump                Dump Pokédex data from a database into CSV files.
-    reindex             Rebuilds the lookup index from the database.
-    setup               Combines load and reindex.
-    status              No effect, but prints which engine, index, and csv
-                        directory would be used for other commands.
-
-Global options:
-    -e|--engine=URI     By default, all commands try to use a SQLite database
-                        in the pokedex install directory.  Use this option (or
-                        a POKEDEX_DB_ENGINE environment variable) to specify an
-                        alternate database.
-    -i|--index=DIR      By default, all commands try to put the lookup index in
-                        the pokedex install directory.  Use this option (or a
-                        POKEDEX_INDEX_DIR environment variable) to specify an
-                        alternate loction.
-    -q|--quiet          Don't print system output.  This is the default for
-                        non-system commands and setup.
-    -v|--verbose        Print system output.  This is the default for system
-                        commands, except setup.
-
-System options:
-    -d|--directory=DIR  By default, load and dump will use the CSV files in the
-                        pokedex install directory.  Use this option to specify
-                        a different directory.
-
-Load options:
-    -D|--drop-tables    Drop all tables before loading data.
-    -S|--safe           Disable engine-specific optimizations.
-    -r|--recursive      Load (and drop) all dependent tables.
-    -l|--langs          Load translations for the given languages.
-                        By default, all available translations are loaded.
-                        Separate multiple languages by a comma (-l en,de,fr)
-
-Dump options:
-    -l|--langs          Dump unofficial texts for given languages.
-                        By default, English (en) is dumped.
-                        Separate multiple languages by a comma (-l en,de,fr)
-                        Use 'none' to not dump any unofficial texts.
-
-    Additionally, load and dump accept a list of table names (possibly with
-    wildcards) and/or csv fileames as an argument list.
-""".encode(locale.getdefaultlocale()[1], 'replace'))
-
-    sys.exit(0)
+def command_help(parser, args):
+    parser.print_help()
 
 
 if __name__ == '__main__':