pokedex/main.py

   1 # encoding: utf8
   2 from optparse import OptionParser
   3 import os
   4 import sys
   5
   6 import pokedex.db
   7 import pokedex.db.load
   8 import pokedex.db.tables
   9 import pokedex.lookup
  10 from pokedex import defaults
  11
  12 def main():
  13     if len(sys.argv) <= 1:
  14         command_help()
  15
  16     command = sys.argv[1]
  17     args = sys.argv[2:]
  18
  19     # XXX there must be a better way to get Unicode argv
  20     # XXX this doesn't work on Windows durp
  21     enc = sys.stdin.encoding or 'utf8'
  22     args = [_.decode(enc) for _ in args]
  23
  24     # Find the command as a function in this file
  25     func = globals().get("command_%s" % command, None)
  26     if func:
  27         func(*args)
  28     else:
  29         command_help()
  30
  31
  32 def get_parser(verbose=True):
  33     """Returns an OptionParser prepopulated with the global options.
  34
  35     `verbose` is whether or not the options should be verbose by default.
  36     """
  37     parser = OptionParser()
  38     parser.add_option('-e', '--engine', dest='engine_uri', default=None)
  39     parser.add_option('-i', '--index', dest='index_dir', default=None)
  40     parser.add_option('-q', '--quiet', dest='verbose', default=verbose, action='store_false')
  41     parser.add_option('-v', '--verbose', dest='verbose', default=verbose, action='store_true')
  42     return parser
  43
  44 def get_session(options):
  45     """Given a parsed options object, connects to the database and returns a
  46     session.
  47     """
  48
  49     engine_uri = options.engine_uri
  50     got_from = 'command line'
  51
  52     if engine_uri is None:
  53         engine_uri, got_from = defaults.get_default_db_uri_with_origin()
  54
  55     session = pokedex.db.connect(engine_uri)
  56
  57     if options.verbose:
  58         print "Connected to database %(engine)s (from %(got_from)s)" \
  59             % dict(engine=session.bind.url, got_from=got_from)
  60
  61     return session
  62
  63 def get_lookup(options, session=None, recreate=False):
  64     """Given a parsed options object, opens the whoosh index and returns a
  65     PokedexLookup object.
  66     """
  67
  68     if recreate and not session:
  69         raise ValueError("get_lookup() needs an explicit session to regen the index")
  70
  71     index_dir = options.index_dir
  72     got_from = 'command line'
  73
  74     if index_dir is None:
  75         index_dir, got_from = defaults.get_default_index_dir_with_origin()
  76
  77     if options.verbose:
  78         print "Opened lookup index %(index_dir)s (from %(got_from)s)" \
  79             % dict(index_dir=index_dir, got_from=got_from)
  80
  81     lookup = pokedex.lookup.PokedexLookup(index_dir, session=session)
  82
  83     if recreate:
  84         lookup.rebuild_index()
  85
  86     return lookup
  87
  88 def get_csv_directory(options):
  89     """Prints and returns the csv directory we're about to use."""
  90
  91     if not options.verbose:
  92         return
  93
  94     csvdir = options.directory
  95     got_from = 'command line'
  96
  97     if csvdir is None:
  98         csvdir, got_from = defaults.get_default_csv_dir_with_origin()
  99
 100     print "Using CSV directory %(csvdir)s (from %(got_from)s)" \
 101         % dict(csvdir=csvdir, got_from=got_from)
 102
 103     return csvdir
 104
 105
 106 ### Plumbing commands
 107
 108 def command_dump(*args):
 109     parser = get_parser(verbose=True)
 110     parser.add_option('-d', '--directory', dest='directory', default=None)
 111     options, tables = parser.parse_args(list(args))
 112
 113     session = get_session(options)
 114     get_csv_directory(options)
 115
 116     pokedex.db.load.dump(session, directory=options.directory,
 117                                   tables=tables,
 118                                   verbose=options.verbose)
 119
 120 def command_load(*args):
 121     parser = get_parser(verbose=True)
 122     parser.add_option('-d', '--directory', dest='directory', default=None)
 123     parser.add_option('-D', '--drop-tables', dest='drop_tables', default=False, action='store_true')
 124     parser.add_option('-S', '--safe', dest='safe', default=False, action='store_true',
 125         help="Do not use backend-specific optimalizations.")
 126     options, tables = parser.parse_args(list(args))
 127
 128     if not options.engine_uri:
 129         print "WARNING: You're reloading the default database, but not the lookup index.  They"
 130         print "         might get out of sync, and pokedex commands may not work correctly!"
 131         print "To fix this, run `pokedex reindex` when this command finishes.  Or, just use"
 132         print "`pokedex setup` to do both at once."
 133         print
 134
 135     session = get_session(options)
 136     get_csv_directory(options)
 137
 138     pokedex.db.load.load(session, directory=options.directory,
 139                                   drop_tables=options.drop_tables,
 140                                   tables=tables,
 141                                   verbose=options.verbose,
 142                                   safe=options.safe)
 143
 144 def command_reindex(*args):
 145     parser = get_parser(verbose=True)
 146     options, _ = parser.parse_args(list(args))
 147
 148     session = get_session(options)
 149     lookup = get_lookup(options, session=session, recreate=True)
 150
 151     print "Recreated lookup index."
 152
 153
 154 def command_setup(*args):
 155     parser = get_parser(verbose=False)
 156     options, _ = parser.parse_args(list(args))
 157
 158     options.directory = None
 159
 160     session = get_session(options)
 161     get_csv_directory(options)
 162     pokedex.db.load.load(session, directory=None, drop_tables=True,
 163                                   verbose=options.verbose,
 164                                   safe=False)
 165
 166     lookup = get_lookup(options, session=session, recreate=True)
 167
 168     print "Recreated lookup index."
 169
 170
 171 def command_status(*args):
 172     parser = get_parser(verbose=True)
 173     options, _ = parser.parse_args(list(args))
 174     options.verbose = True
 175     options.directory = None
 176
 177     # Database, and a lame check for whether it's been inited at least once
 178     session = get_session(options)
 179     print "  - OK!  Connected successfully."
 180
 181     if pokedex.db.tables.Pokemon.__table__.exists(session.bind):
 182         print "  - OK!  Database seems to contain some data."
 183     else:
 184         print "  - WARNING: Database appears to be empty."
 185
 186     # CSV; simple checks that the dir exists
 187     csvdir = get_csv_directory(options)
 188     if not os.path.exists(csvdir):
 189         print "  - ERROR: No such directory!"
 190     elif not os.path.isdir(csvdir):
 191         print "  - ERROR: Not a directory!"
 192     else:
 193         print "  - OK!  Directory exists."
 194
 195         if os.access(csvdir, os.R_OK):
 196             print "  - OK!  Can read from directory."
 197         else:
 198             print "  - ERROR: Can't read from directory!"
 199
 200         if os.access(csvdir, os.W_OK):
 201             print "  - OK!  Can write to directory."
 202         else:
 203             print "  - WARNING: Can't write to directory!  " \
 204                 "`dump` will not work.  You may need to sudo."
 205
 206     # Index; the PokedexLookup constructor covers most tests and will
 207     # cheerfully bomb if they fail
 208     lookup = get_lookup(options, recreate=False)
 209     print "  - OK!  Opened successfully."
 210
 211
 212 ### User-facing commands
 213
 214 def command_lookup(*args):
 215     parser = get_parser(verbose=False)
 216     options, words = parser.parse_args(list(args))
 217
 218     name = u' '.join(words)
 219
 220     session = get_session(options)
 221     lookup = get_lookup(options, session=session, recreate=False)
 222
 223     results = lookup.lookup(name)
 224     if not results:
 225         print "No matches."
 226     elif results[0].exact:
 227         print "Matched:"
 228     else:
 229         print "Fuzzy-matched:"
 230
 231     for result in results:
 232         if hasattr(result.object, 'full_name'):
 233             name = result.object.full_name
 234         else:
 235             name = result.object.name
 236
 237         print "%s: %s" % (result.object.__tablename__, name),
 238         if result.language:
 239             print "(%s in %s)" % (result.name, result.language)
 240         else:
 241             print
 242
 243
 244 def command_help():
 245     print u"""pokedex -- a command-line Pokédex interface
 246 usage: pokedex {command} [options...]
 247 Run `pokedex setup` first, or nothing will work!
 248 See http://bugs.veekun.com/projects/pokedex/wiki/CLI for more documentation.
 249
 250 Commands:
 251     help                Displays this message.
 252     lookup [thing]      Look up something in the Pokédex.
 253
 254 System commands:
 255     load                Load Pokédex data into a database from CSV files.
 256     dump                Dump Pokédex data from a database into CSV files.
 257     reindex             Rebuilds the lookup index from the database.
 258     setup               Combines load and reindex.
 259     status              No effect, but prints which engine, index, and csv
 260                         directory would be used for other commands.
 261
 262 Global options:
 263     -e|--engine=URI     By default, all commands try to use a SQLite database
 264                         in the pokedex install directory.  Use this option (or
 265                         a POKEDEX_DB_ENGINE environment variable) to specify an
 266                         alternate database.
 267     -i|--index=DIR      By default, all commands try to put the lookup index in
 268                         the pokedex install directory.  Use this option (or a
 269                         POKEDEX_INDEX_DIR environment variable) to specify an
 270                         alternate loction.
 271     -q|--quiet          Don't print system output.  This is the default for
 272                         non-system commands and setup.
 273     -v|--verbose        Print system output.  This is the default for system
 274                         commands, except setup.
 275
 276 System options:
 277     -d|--directory=DIR  By default, load and dump will use the CSV files in the
 278                         pokedex install directory.  Use this option to specify
 279                         a different directory.
 280     -D|--drop-tables    With load, drop all tables before loading data.
 281
 282     Additionally, load and dump accept a list of table names (possibly with
 283     wildcards) and/or csv fileames as an argument list.
 284 """.encode(sys.getdefaultencoding(), 'replace')
 285
 286     sys.exit(0)