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