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