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