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