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