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