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