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('-r', '--recursive', dest='recursive', default=False, action='store_true')
 125     parser.add_option('-S', '--safe', dest='safe', default=False, action='store_true',
 126         help="Do not use backend-specific optimalizations.")
 127     options, tables = parser.parse_args(list(args))
 128
 129     if not options.engine_uri:
 130         print "WARNING: You're reloading the default database, but not the lookup index.  They"
 131         print "         might get out of sync, and pokedex commands may not work correctly!"
 132         print "To fix this, run `pokedex reindex` when this command finishes.  Or, just use"
 133         print "`pokedex setup` to do both at once."
 134         print
 135
 136     session = get_session(options)
 137     get_csv_directory(options)
 138
 139     pokedex.db.load.load(session, directory=options.directory,
 140                                   drop_tables=options.drop_tables,
 141                                   tables=tables,
 142                                   verbose=options.verbose,
 143                                   safe=options.safe,
 144                                   recursive=options.recursive)
 145
 146 def command_reindex(*args):
 147     parser = get_parser(verbose=True)
 148     options, _ = parser.parse_args(list(args))
 149
 150     session = get_session(options)
 151     lookup = get_lookup(options, session=session, recreate=True)
 152
 153     print "Recreated lookup index."
 154
 155
 156 def command_setup(*args):
 157     parser = get_parser(verbose=False)
 158     options, _ = parser.parse_args(list(args))
 159
 160     options.directory = None
 161
 162     session = get_session(options)
 163     get_csv_directory(options)
 164     pokedex.db.load.load(session, directory=None, drop_tables=True,
 165                                   verbose=options.verbose,
 166                                   safe=False)
 167
 168     lookup = get_lookup(options, session=session, recreate=True)
 169
 170     print "Recreated lookup index."
 171
 172
 173 def command_status(*args):
 174     parser = get_parser(verbose=True)
 175     options, _ = parser.parse_args(list(args))
 176     options.verbose = True
 177     options.directory = None
 178
 179     # Database, and a lame check for whether it's been inited at least once
 180     session = get_session(options)
 181     print "  - OK!  Connected successfully."
 182
 183     if pokedex.db.tables.Pokemon.__table__.exists(session.bind):
 184         print "  - OK!  Database seems to contain some data."
 185     else:
 186         print "  - WARNING: Database appears to be empty."
 187
 188     # CSV; simple checks that the dir exists
 189     csvdir = get_csv_directory(options)
 190     if not os.path.exists(csvdir):
 191         print "  - ERROR: No such directory!"
 192     elif not os.path.isdir(csvdir):
 193         print "  - ERROR: Not a directory!"
 194     else:
 195         print "  - OK!  Directory exists."
 196
 197         if os.access(csvdir, os.R_OK):
 198             print "  - OK!  Can read from directory."
 199         else:
 200             print "  - ERROR: Can't read from directory!"
 201
 202         if os.access(csvdir, os.W_OK):
 203             print "  - OK!  Can write to directory."
 204         else:
 205             print "  - WARNING: Can't write to directory!  " \
 206                 "`dump` will not work.  You may need to sudo."
 207
 208     # Index; the PokedexLookup constructor covers most tests and will
 209     # cheerfully bomb if they fail
 210     lookup = get_lookup(options, recreate=False)
 211     print "  - OK!  Opened successfully."
 212
 213
 214 ### User-facing commands
 215
 216 def command_lookup(*args):
 217     parser = get_parser(verbose=False)
 218     options, words = parser.parse_args(list(args))
 219
 220     name = u' '.join(words)
 221
 222     session = get_session(options)
 223     lookup = get_lookup(options, session=session, recreate=False)
 224
 225     results = lookup.lookup(name)
 226     if not results:
 227         print "No matches."
 228     elif results[0].exact:
 229         print "Matched:"
 230     else:
 231         print "Fuzzy-matched:"
 232
 233     for result in results:
 234         if hasattr(result.object, 'full_name'):
 235             name = result.object.full_name
 236         else:
 237             name = result.object.name
 238
 239         print "%s: %s" % (result.object.__tablename__, name),
 240         if result.language:
 241             print "(%s in %s)" % (result.name, result.language)
 242         else:
 243             print
 244
 245
 246 def command_help():
 247     print u"""pokedex -- a command-line Pokédex interface
 248 usage: pokedex {command} [options...]
 249 Run `pokedex setup` first, or nothing will work!
 250 See http://bugs.veekun.com/projects/pokedex/wiki/CLI for more documentation.
 251
 252 Commands:
 253     help                Displays this message.
 254     lookup [thing]      Look up something in the Pokédex.
 255
 256 System commands:
 257     load                Load Pokédex data into a database from CSV files.
 258     dump                Dump Pokédex data from a database into CSV files.
 259     reindex             Rebuilds the lookup index from the database.
 260     setup               Combines load and reindex.
 261     status              No effect, but prints which engine, index, and csv
 262                         directory would be used for other commands.
 263
 264 Global options:
 265     -e|--engine=URI     By default, all commands try to use a SQLite database
 266                         in the pokedex install directory.  Use this option (or
 267                         a POKEDEX_DB_ENGINE environment variable) to specify an
 268                         alternate database.
 269     -i|--index=DIR      By default, all commands try to put the lookup index in
 270                         the pokedex install directory.  Use this option (or a
 271                         POKEDEX_INDEX_DIR environment variable) to specify an
 272                         alternate loction.
 273     -q|--quiet          Don't print system output.  This is the default for
 274                         non-system commands and setup.
 275     -v|--verbose        Print system output.  This is the default for system
 276                         commands, except setup.
 277
 278 System options:
 279     -d|--directory=DIR  By default, load and dump will use the CSV files in the
 280                         pokedex install directory.  Use this option to specify
 281                         a different directory.
 282
 283 Load options:
 284     -D|--drop-tables    Drop all tables before loading data.
 285     -S|--safe           Disable engine-specific optimizations.
 286     -r|--recursive      Load (and drop) all dependent tables.
 287
 288     Additionally, load and dump accept a list of table names (possibly with
 289     wildcards) and/or csv fileames as an argument list.
 290 """.encode(sys.getdefaultencoding(), 'replace')
 291
 292     sys.exit(0)