veekun_pokedex/pokedex/main.py

346 lines
11 KiB
Python
Raw Normal View History

# encoding: utf8
2015-10-05 15:11:08 +00:00
from __future__ import print_function
2016-01-28 12:01:23 +00:00
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
2016-01-28 12:01:23 +00:00
def main(junk, *argv):
parser = create_parser()
2016-01-28 12:01:23 +00:00
if len(argv) <= 0:
parser.print_help()
sys.exit()
2016-01-28 12:01:23 +00:00
args = parser.parse_args(argv)
args.func(parser, args)
def setuptools_entry():
main(*sys.argv)
2016-01-28 12:01:23 +00:00
def create_parser():
"""Build and return an ArgumentParser.
"""
# Slightly clumsy workaround to make both `setup -v` and `-v setup` work
common_parser = argparse.ArgumentParser(add_help=False)
common_parser.add_argument(
2016-01-28 12:01:23 +00:00
'-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.',
)
common_parser.add_argument(
2016-01-28 12:01:23 +00:00
'-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.',
)
common_parser.add_argument(
2016-01-28 12:01:23 +00:00
'-q', '--quiet', dest='verbose', action='store_false',
help=u'Don\'t print system output. This is the default for '
'non-system commands and setup.',
)
common_parser.add_argument(
2016-01-28 12:01:23 +00:00
'-v', '--verbose', dest='verbose', default=False, action='store_true',
help=u'Print system output. This is the default for system '
u'commands, except setup.',
)
parser = argparse.ArgumentParser(
prog='pokedex', description=u'A command-line Pokédex interface',
parents=[common_parser],
)
cmds = parser.add_subparsers(title='commands', metavar='<command>', help='commands')
cmd_help = cmds.add_parser(
'help', help=u'Display this message',
parents=[common_parser])
2016-01-28 12:01:23 +00:00
cmd_help.set_defaults(func=command_help)
cmd_lookup = cmds.add_parser(
'lookup', help=u'Look up something in the Pokédex',
parents=[common_parser])
2016-01-28 12:01:23 +00:00
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',
parents=[common_parser])
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',
parents=[common_parser])
2016-01-28 12:01:23 +00:00
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',
parents=[common_parser])
2016-01-28 12:01:23 +00:00
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',
parents=[common_parser])
2016-01-28 12:01:23 +00:00
cmd_reindex.set_defaults(func=command_reindex, verbose=True)
cmd_setup = cmds.add_parser(
'setup', help=u'Combine load and reindex',
parents=[common_parser])
2016-01-28 12:01:23 +00:00
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',
parents=[common_parser])
2016-01-28 12:01:23 +00:00
cmd_status.set_defaults(func=command_status, verbose=True)
return parser
2016-01-28 12:01:23 +00:00
def get_session(args):
"""Given a parsed options object, connects to the database and returns a
session.
"""
2016-01-28 12:01:23 +00:00
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)
2016-01-28 12:01:23 +00:00
if args.verbose:
2015-10-05 15:11:08 +00:00
print("Connected to database %(engine)s (from %(got_from)s)"
% dict(engine=session.bind.url, got_from=got_from))
return session
2016-01-28 12:01:23 +00:00
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")
2016-01-28 12:01:23 +00:00
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()
2016-01-28 12:01:23 +00:00
if args.verbose:
2015-10-05 15:11:08 +00:00
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
2016-01-28 12:01:23 +00:00
def get_csv_directory(args):
"""Prints and returns the csv directory we're about to use."""
2016-01-28 12:01:23 +00:00
if not args.verbose:
return
2016-01-28 12:01:23 +00:00
csvdir = args.directory
got_from = 'command line'
if csvdir is None:
csvdir, got_from = defaults.get_default_csv_dir_with_origin()
2015-10-05 15:11:08 +00:00
print("Using CSV directory %(csvdir)s (from %(got_from)s)"
% dict(csvdir=csvdir, got_from=got_from))
return csvdir
### Plumbing commands
2016-01-28 12:01:23 +00:00
def command_dump(parser, args):
session = get_session(args)
get_csv_directory(args)
2016-01-28 12:01:23 +00:00
if args.langs is not None:
langs = [l.strip() for l in args.langs.split(',')]
else:
langs = None
2016-01-28 12:01:23 +00:00
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:
2015-10-05 15:11:08 +00:00
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()
2016-01-28 12:01:23 +00:00
if args.langs == 'none':
2011-04-03 17:05:21 +00:00
langs = []
2016-01-28 12:01:23 +00:00
elif args.langs is None:
2011-04-03 17:05:21 +00:00
langs = None
else:
2016-01-28 12:01:23 +00:00
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)
2015-10-05 15:11:08 +00:00
print("Recreated lookup index.")
2016-01-28 12:01:23 +00:00
def command_setup(parser, args):
args.directory = None
2016-01-28 12:01:23 +00:00
session = get_session(args)
get_csv_directory(args)
pokedex.db.load.load(
session, directory=None, drop_tables=True,
verbose=args.verbose, safe=False)
2016-01-28 12:01:23 +00:00
get_lookup(args, session=session, recreate=True)
2015-10-05 15:11:08 +00:00
print("Recreated lookup index.")
2016-01-28 12:01:23 +00:00
def command_status(parser, args):
args.directory = None
# Database, and a lame check for whether it's been inited at least once
2016-01-28 12:01:23 +00:00
session = get_session(args)
2015-10-05 15:11:08 +00:00
print(" - OK! Connected successfully.")
if pokedex.db.tables.Pokemon.__table__.exists(session.bind):
2015-10-05 15:11:08 +00:00
print(" - OK! Database seems to contain some data.")
else:
2015-10-05 15:11:08 +00:00
print(" - WARNING: Database appears to be empty.")
# CSV; simple checks that the dir exists
2016-01-28 12:01:23 +00:00
csvdir = get_csv_directory(args)
if not os.path.exists(csvdir):
2015-10-05 15:11:08 +00:00
print(" - ERROR: No such directory!")
elif not os.path.isdir(csvdir):
2015-10-05 15:11:08 +00:00
print(" - ERROR: Not a directory!")
else:
2015-10-05 15:11:08 +00:00
print(" - OK! Directory exists.")
if os.access(csvdir, os.R_OK):
2015-10-05 15:11:08 +00:00
print(" - OK! Can read from directory.")
else:
2015-10-05 15:11:08 +00:00
print(" - ERROR: Can't read from directory!")
if os.access(csvdir, os.W_OK):
2015-10-05 15:11:08 +00:00
print(" - OK! Can write to directory.")
else:
2015-10-05 15:11:08 +00:00
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
2016-01-28 12:01:23 +00:00
get_lookup(args, recreate=False)
2015-10-05 15:11:08 +00:00
print(" - OK! Opened successfully.")
### User-facing commands
2016-01-28 12:01:23 +00:00
def command_lookup(parser, args):
name = u' '.join(args.criteria)
2016-01-28 12:01:23 +00:00
session = get_session(args)
lookup = get_lookup(args, session=session, recreate=False)
results = lookup.lookup(name)
if not results:
2015-10-05 15:11:08 +00:00
print("No matches.")
elif results[0].exact:
2015-10-05 15:11:08 +00:00
print("Matched:")
else:
2015-10-05 15:11:08 +00:00
print("Fuzzy-matched:")
for result in results:
if hasattr(result.object, 'full_name'):
name = result.object.full_name
else:
name = result.object.name
2015-10-05 15:11:08 +00:00
print("%s: %s" % (result.object.__tablename__, name), end='')
if result.language:
2015-10-05 15:11:08 +00:00
print("(%s in %s)" % (result.name, result.language))
else:
2015-10-05 15:11:08 +00:00
print()
2016-01-28 12:01:23 +00:00
def command_help(parser, args):
parser.print_help()
if __name__ == '__main__':
main(*sys.argv)