veekun_pokedex/pokedex/main.py

# encoding: utf8
from __future__ import print_function

import argparse
import os
import sys

import pokedex.cli.search
import pokedex.db
import pokedex.db.load
import pokedex.db.tables
import pokedex.lookup
from pokedex import defaults


def main(junk, *argv):
    if len(argv) <= 0:
        command_help()
        return

    parser = create_parser()
    args = parser.parse_args(argv)
    args.func(parser, args)


def setuptools_entry():
    main(*sys.argv)


def create_parser():
    """Build and return an ArgumentParser.
    """
    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_search = cmds.add_parser('search', help=u'Find things by various criteria')
    pokedex.cli.search.configure_parser(cmd_search)

    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(args):
    """Given a parsed options object, connects to the database and returns a
    session.
    """

    engine_uri = args.engine_uri
    got_from = 'command line'

    if engine_uri is None:
        engine_uri, got_from = defaults.get_default_db_uri_with_origin()

    session = pokedex.db.connect(engine_uri)

    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(args, session=None, recreate=False):
    """Given a parsed options object, opens the whoosh index and returns a
    PokedexLookup object.
    """

    if recreate and not session:
        raise ValueError("get_lookup() needs an explicit session to regen the index")

    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 args.verbose:
        print("Opened lookup index %(index_dir)s (from %(got_from)s)"
            % dict(index_dir=index_dir, got_from=got_from))

    lookup = pokedex.lookup.PokedexLookup(index_dir, session=session)

    if recreate:
        lookup.rebuild_index()

    return lookup


def get_csv_directory(args):
    """Prints and returns the csv directory we're about to use."""

    if not args.verbose:
        return

    csvdir = args.directory
    got_from = 'command line'

    if csvdir is None:
        csvdir, got_from = defaults.get_default_csv_dir_with_origin()

    print("Using CSV directory %(csvdir)s (from %(got_from)s)"
        % dict(csvdir=csvdir, got_from=got_from))

    return csvdir


### Plumbing commands

def command_dump(parser, args):
    session = get_session(args)
    get_csv_directory(args)

    if args.langs is not None:
        langs = [l.strip() for l in args.langs.split(',')]
    else:
        langs = None

    pokedex.db.load.dump(
        session,
        directory=args.directory,
        tables=args.tables,
        verbose=args.verbose,
        langs=langs,
    )


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 args.langs == 'none':
        langs = []
    elif args.langs is None:
        langs = None
    else:
        langs = [l.strip() for l in args.langs.split(',')]

    session = get_session(args)
    get_csv_directory(args)

    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(parser, args):
    session = get_session(args)
    get_lookup(args, session=session, recreate=True)
    print("Recreated lookup index.")


def command_setup(parser, args):
    args.directory = None

    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(parser, args):
    args.directory = None

    # Database, and a lame check for whether it's been inited at least once
    session = get_session(args)
    print("  - OK!  Connected successfully.")

    if pokedex.db.tables.Pokemon.__table__.exists(session.bind):
        print("  - OK!  Database seems to contain some data.")
    else:
        print("  - WARNING: Database appears to be empty.")

    # CSV; simple checks that the dir exists
    csvdir = get_csv_directory(args)
    if not os.path.exists(csvdir):
        print("  - ERROR: No such directory!")
    elif not os.path.isdir(csvdir):
        print("  - ERROR: Not a directory!")
    else:
        print("  - OK!  Directory exists.")

        if os.access(csvdir, os.R_OK):
            print("  - OK!  Can read from directory.")
        else:
            print("  - ERROR: Can't read from directory!")

        if os.access(csvdir, os.W_OK):
            print("  - OK!  Can write to directory.")
        else:
            print("  - WARNING: Can't write to directory!  "
                "`dump` will not work.  You may need to sudo.")

    # Index; the PokedexLookup constructor covers most tests and will
    # cheerfully bomb if they fail
    get_lookup(args, recreate=False)
    print("  - OK!  Opened successfully.")


### User-facing commands

def command_lookup(parser, args):
    name = u' '.join(args.criteria)

    session = get_session(args)
    lookup = get_lookup(args, session=session, recreate=False)

    results = lookup.lookup(name)
    if not results:
        print("No matches.")
    elif results[0].exact:
        print("Matched:")
    else:
        print("Fuzzy-matched:")

    for result in results:
        if hasattr(result.object, 'full_name'):
            name = result.object.full_name
        else:
            name = result.object.name

        print("%s: %s" % (result.object.__tablename__, name), end='')
        if result.language:
            print("(%s in %s)" % (result.name, result.language))
        else:
            print()


def command_help(parser, args):
    parser.print_help()


if __name__ == '__main__':
    main(*sys.argv)