810783f2ba72a2c67b2d476f66c5aa17cd28764e
[zzz-pokedex.git] / 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)