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
  98     if recreate:
  99         lookup.rebuild_index()
 100
 101     return lookup
 102
 103 def print_csv_directory(options):
 104     """Just prints the csv directory we're about to use."""
 105
 106     if not options.verbose:
 107         return
 108
 109     if options.directory:
 110         csvdir = options.directory
 111         got_from = 'command line'
 112     else:
 113         # This is the same as the db.load default
 114         csvdir = pkg_resources.resource_filename('pokedex', 'data/csv')
 115         got_from = 'default setting'
 116
 117     print "Using CSV directory {csvdir} (from {got_from})" \
 118         .format(csvdir=csvdir, got_from=got_from)
 119
 120
 121 ### Plumbing commands
 122
 123 def command_dump(*args):
 124     parser = get_parser(verbose=True)
 125     parser.add_option('-d', '--directory', dest='directory', default=None)
 126     options, tables = parser.parse_args(list(args))
 127
 128     session = get_session(options)
 129     print_csv_directory(options)
 130
 131     pokedex.db.load.dump(session, directory=options.directory,
 132                                   tables=tables,
 133                                   verbose=options.verbose)
 134
 135 def command_load(*args):
 136     parser = get_parser(verbose=True)
 137     parser.add_option('-d', '--directory', dest='directory', default=None)
 138     parser.add_option('-D', '--drop-tables', dest='drop_tables', default=False, action='store_true')
 139     options, tables = parser.parse_args(list(args))
 140
 141     if not options.engine_uri:
 142         print "WARNING: You're reloading the default database, but not the lookup index.  They"
 143         print "         might get out of sync, and pokedex commands may not work correctly!"
 144         print "To fix this, run `pokedex reindex` when this command finishes.  Or, just use"
 145         print "`pokedex setup` to do both at once."
 146         print
 147
 148     session = get_session(options)
 149     print_csv_directory(options)
 150
 151     pokedex.db.load.load(session, directory=options.directory,
 152                                   drop_tables=options.drop_tables,
 153                                   tables=tables,
 154                                   verbose=options.verbose)
 155
 156 def command_reindex(*args):
 157     parser = get_parser(verbose=True)
 158     options, _ = parser.parse_args(list(args))
 159
 160     session = get_session(options)
 161     lookup = get_lookup(options, session=session, recreate=True)
 162
 163     print "Recreated lookup index."
 164
 165
 166 def command_setup(*args):
 167     parser = get_parser(verbose=False)
 168     options, _ = parser.parse_args(list(args))
 169
 170     options.directory = None
 171
 172     session = get_session(options)
 173     print_csv_directory(options)
 174     pokedex.db.load.load(session, directory=None, drop_tables=True,
 175                                   verbose=options.verbose)
 176
 177     lookup = get_lookup(options, session=session, recreate=True)
 178
 179     print "Recreated lookup index."
 180
 181
 182 def command_status(*args):
 183     parser = get_parser(verbose=True)
 184     options, _ = parser.parse_args(list(args))
 185     options.verbose = True
 186     options.directory = None
 187
 188     session = get_session(options)
 189     print_csv_directory(options)
 190     lookup = get_lookup(options, recreate=False)
 191
 192
 193 ### User-facing commands
 194
 195 def command_lookup(*args):
 196     parser = get_parser(verbose=False)
 197     options, words = parser.parse_args(list(args))
 198
 199     name = u' '.join(words)
 200
 201     session = get_session(options)
 202     lookup = get_lookup(options, session=session, recreate=False)
 203
 204     results = lookup.lookup(name)
 205     if not results:
 206         print "No matches."
 207     elif results[0].exact:
 208         print "Matched:"
 209     else:
 210         print "Fuzzy-matched:"
 211
 212     for result in results:
 213         if hasattr(result.object, 'full_name'):
 214             name = result.object.full_name
 215         else:
 216             name = result.object.name
 217
 218         print "%s: %s" % (result.object.__tablename__, name),
 219         if result.language:
 220             print "(%s in %s)" % (result.name, result.language)
 221         else:
 222             print
 223
 224
 225 def command_help():
 226     print u"""pokedex -- a command-line Pokédex interface
 227 usage: pokedex {command} [options...]
 228 Run `pokedex setup` first, or nothing will work!
 229 See http://bugs.veekun.com/projects/pokedex/wiki/CLI for more documentation.
 230
 231 Commands:
 232     help                Displays this message.
 233     lookup [thing]      Look up something in the Pokédex.
 234
 235 System commands:
 236     load                Load Pokédex data into a database from CSV files.
 237     dump                Dump Pokédex data from a database into CSV files.
 238     reindex             Rebuilds the lookup index from the database.
 239     setup               Combines load and reindex.
 240     status              No effect, but prints which engine, index, and csv
 241                         directory would be used for other commands.
 242
 243 Global options:
 244     -e|--engine=URI     By default, all commands try to use a SQLite database
 245                         in the pokedex install directory.  Use this option (or
 246                         a POKEDEX_DB_ENGINE environment variable) to specify an
 247                         alternate database.
 248     -i|--index=DIR      By default, all commands try to put the lookup index in
 249                         the pokedex install directory.  Use this option (or a
 250                         POKEDEX_INDEX_DIR environment variable) to specify an
 251                         alternate loction.
 252     -q|--quiet          Don't print system output.  This is the default for
 253                         non-system commands and setup.
 254     -v|--verbose        Print system output.  This is the default for system
 255                         commands, except setup.
 256
 257 System options:
 258     -d|--directory=DIR  By default, load and dump will use the CSV files in the
 259                         pokedex install directory.  Use this option to specify
 260                         a different directory.
 261     -D|--drop-tables    With load, drop all tables before loading data.
 262
 263     Additionally, load and dump accept a list of table names (possibly with
 264     wildcards) and/or csv fileames as an argument list.
 265 """.encode(sys.getdefaultencoding(), 'replace')
 266
 267     sys.exit(0)