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