pokedex/__init__.py

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