f35e335190275f92a9e27026af581e9db6ac89ae
[zzz-pokedex.git] / pokedex / lookup.py
1 # encoding: utf8
2 from collections import namedtuple
3 import os, os.path
4 import pkg_resources
5 import random
6 import re
7 import shutil
8 import unicodedata
9
10 from sqlalchemy.sql import func
11 import whoosh
12 import whoosh.filedb.filestore
13 import whoosh.filedb.fileindex
14 import whoosh.index
15 from whoosh.qparser import QueryParser
16 import whoosh.scoring
17 import whoosh.spelling
18
19 from pokedex.db import connect
20 import pokedex.db.tables as tables
21 from pokedex.roomaji import romanize
22
23 __all__ = ['PokedexLookup']
24
25
26 rx_is_number = re.compile('^\d+$')
27
28 LookupResult = namedtuple('LookupResult',
29 ['object', 'indexed_name', 'name', 'language', 'iso3166', 'exact'])
30
31 class LanguageWeighting(whoosh.scoring.Weighting):
32 """A scoring class that forces otherwise-equal English results to come
33 before foreign results.
34 """
35
36 def score(self, searcher, fieldnum, text, docnum, weight, QTF=1):
37 doc = searcher.stored_fields(docnum)
38 if doc['language'] == None:
39 # English (well, "default"); leave it at 1
40 return weight
41 elif doc['language'] == u'Roomaji':
42 # Give Roomaji a little boost; it's most likely to be searched
43 return weight * 0.95
44 else:
45 # Everything else can drop down the totem pole
46 return weight * 0.9
47
48
49 class PokedexLookup(object):
50 INTERMEDIATE_LOOKUP_RESULTS = 25
51 MAX_LOOKUP_RESULTS = 10
52
53 # Dictionary of table name => table class.
54 # Need the table name so we can get the class from the table name after we
55 # retrieve something from the index
56 indexed_tables = dict(
57 (cls.__tablename__, cls)
58 for cls in (
59 tables.Ability,
60 tables.Item,
61 tables.Move,
62 tables.Pokemon,
63 tables.Type,
64 )
65 )
66
67
68 def __init__(self, directory=None, session=None, recreate=False):
69 """Opens the whoosh index stored in the named directory. If the index
70 doesn't already exist, it will be created.
71
72 `directory`
73 Directory containing the index. Defaults to a location within the
74 `pokedex` egg directory.
75
76 `session`
77 If the index needs to be created, this database session will be
78 used. Defaults to an attempt to connect to the default SQLite
79 database installed by `pokedex setup`.
80
81 `recreate`
82 If set to True, the whoosh index will be created even if it already
83 exists.
84 """
85
86 # By the time this returns, self.index, self.speller, and self.session
87 # must be set
88
89 # Defaults
90 if not directory:
91 directory = pkg_resources.resource_filename('pokedex',
92 'data/whoosh-index')
93
94 if session:
95 self.session = session
96 else:
97 self.session = connect()
98
99 # Attempt to open or create the index
100 directory_exists = os.path.exists(directory)
101 if directory_exists and not recreate:
102 # Already exists; should be an index! Bam, done.
103 try:
104 self.index = whoosh.index.open_dir(directory, indexname='MAIN')
105 spell_store = whoosh.filedb.filestore.FileStorage(directory)
106 self.speller = whoosh.spelling.SpellChecker(spell_store)
107 return
108 except whoosh.index.EmptyIndexError as e:
109 # Apparently not a real index. Fall out and create it
110 pass
111
112 # Delete and start over if we're going to bail anyway.
113 if directory_exists and recreate:
114 # Be safe and only delete if it looks like a whoosh index, i.e.,
115 # everything starts with _
116 if all(f[0] == '_' for f in os.listdir(directory)):
117 shutil.rmtree(directory)
118 directory_exists = False
119
120 if not directory_exists:
121 os.mkdir(directory)
122
123
124 ### Create index
125 schema = whoosh.fields.Schema(
126 name=whoosh.fields.ID(stored=True),
127 table=whoosh.fields.ID(stored=True),
128 row_id=whoosh.fields.ID(stored=True),
129 language=whoosh.fields.STORED,
130 iso3166=whoosh.fields.STORED,
131 display_name=whoosh.fields.STORED, # non-lowercased name
132 )
133
134 self.index = whoosh.index.create_in(directory, schema=schema,
135 indexname='MAIN')
136 writer = self.index.writer()
137
138 # Index every name in all our tables of interest
139 # speller_entries becomes a list of (word, score) tuples; the score is
140 # 2 for English names, 1.5 for Roomaji, and 1 for everything else. I
141 # think this biases the results in the direction most people expect,
142 # especially when e.g. German names are very similar to English names
143 speller_entries = []
144 for cls in self.indexed_tables.values():
145 q = session.query(cls)
146
147 for row in q.yield_per(5):
148 row_key = dict(table=unicode(cls.__tablename__),
149 row_id=unicode(row.id))
150
151 def add(name, language, iso3166, score):
152 normalized_name = self.normalize_name(name)
153
154 writer.add_document(
155 name=normalized_name, display_name=name,
156 language=language, iso3166=iso3166,
157 **row_key
158 )
159
160 speller_entries.append((normalized_name, score))
161
162
163 # Add the basic English name to the index
164 if cls == tables.Pokemon:
165 # Pokémon need their form name added
166 # XXX kinda kludgy
167 add(row.full_name, None, u'us', 1)
168
169 # If this is a default form, ALSO add the unadorned name,
170 # so 'Deoxys' alone will still do the right thing
171 if row.forme_name and not row.forme_base_pokemon_id:
172 add(row.name, None, u'us', 1)
173 else:
174 add(row.name, None, u'us', 1)
175
176 # Some things also have other languages' names
177 # XXX other language form names..?
178 for foreign_name in getattr(row, 'foreign_names', []):
179 moonspeak = foreign_name.name
180 if row.name == moonspeak:
181 # Don't add the English name again as a different
182 # language; no point and it makes spell results
183 # confusing
184 continue
185
186 add(moonspeak, foreign_name.language.name,
187 foreign_name.language.iso3166,
188 3)
189
190 # Add Roomaji too
191 if foreign_name.language.name == 'Japanese':
192 roomaji = romanize(foreign_name.name)
193 add(roomaji, u'Roomaji', u'jp', 8)
194
195 writer.commit()
196
197 # Construct and populate a spell-checker index. Quicker to do it all
198 # at once, as every call to add_* does a commit(), and those seem to be
199 # expensive
200 self.speller = whoosh.spelling.SpellChecker(self.index.storage)
201 self.speller.add_scored_words(speller_entries)
202
203
204 def normalize_name(self, name):
205 """Strips irrelevant formatting junk from name input.
206
207 Specifically: everything is lowercased, and accents are removed.
208 """
209 # http://stackoverflow.com/questions/517923/what-is-the-best-way-to-remove-accents-in-a-python-unicode-string
210 # Makes sense to me. Decompose by Unicode rules, then remove combining
211 # characters, then recombine. I'm explicitly doing it this way instead
212 # of testing combining() because Korean characters apparently
213 # decompose! But the results are considered letters, not combining
214 # characters, so testing for Mn works well, and combining them again
215 # makes them look right.
216 nkfd_form = unicodedata.normalize('NFKD', unicode(name))
217 name = u"".join(c for c in nkfd_form
218 if unicodedata.category(c) != 'Mn')
219 name = unicodedata.normalize('NFC', name)
220
221 name = name.strip()
222 name = name.lower()
223
224 return name
225
226
227 def _parse_table_name(self, name):
228 """Takes a singular table name, table name, or table object and returns
229 the table name.
230
231 Returns None for a bogus name.
232 """
233 if hasattr(name, '__tablename__'):
234 return getattr(name, '__tablename__')
235 elif name in self.indexed_tables:
236 return name
237 elif name + 's' in self.indexed_tables:
238 return name + 's'
239 else:
240 # Bogus. Be nice and return dummy
241 return None
242
243 def _whoosh_records_to_results(self, records, exact=True):
244 """Converts a list of whoosh's indexed records to LookupResult tuples
245 containing database objects.
246 """
247 # XXX this 'exact' thing is getting kinda leaky. would like a better
248 # way to handle it, since only lookup() cares about fuzzy results
249 seen = {}
250 results = []
251 for record in records:
252 # Skip dupes
253 seen_key = record['table'], record['row_id']
254 if seen_key in seen:
255 continue
256 seen[seen_key] = True
257
258 cls = self.indexed_tables[record['table']]
259 obj = self.session.query(cls).get(record['row_id'])
260
261 results.append(LookupResult(object=obj,
262 indexed_name=record['name'],
263 name=record['display_name'],
264 language=record['language'],
265 iso3166=record['iso3166'],
266 exact=exact))
267
268 return results
269
270
271 def lookup(self, input, valid_types=[], exact_only=False):
272 """Attempts to find some sort of object, given a name.
273
274 Returns a list of named (object, name, language, iso3166, exact)
275 tuples. `object` is a database object, `name` is the name under which
276 the object was found, `language` and `iso3166` are the name and country
277 code of the language in which the name was found, and `exact` is True
278 iff this was an
279 exact match.
280
281 This function currently ONLY does fuzzy matching if there are no exact
282 matches.
283
284 Formes are not returned unless requested; "Shaymin" will return only
285 grass Shaymin.
286
287 Extraneous whitespace is removed with extreme prejudice.
288
289 Recognizes:
290 - Names: "Eevee", "Surf", "Run Away", "Payapa Berry", etc.
291 - Foreign names: "Iibui", "Eivui"
292 - Fuzzy names in whatever language: "Evee", "Ibui"
293 - IDs: "133", "192", "250"
294 Also:
295 - Type restrictions. "type:psychic" will only return the type. This
296 is how to make ID lookup useful. Multiple type specs can be entered
297 with commas, as "move,item:1". If `valid_types` are provided, any
298 type prefix will be ignored.
299 - Alternate formes can be specified merely like "wash rotom".
300
301 `input`
302 Name of the thing to look for.
303
304 `valid_types`
305 A list of table objects or names, e.g., `['pokemon', 'moves']`. If
306 this is provided, only results in one of the given tables will be
307 returned.
308
309 `exact_only`
310 If True, only exact matches are returned. If set to False (the
311 default), and the provided `name` doesn't match anything exactly,
312 spelling correction will be attempted.
313 """
314
315 name = self.normalize_name(input)
316 exact = True
317 form = None
318
319 # Remove any type prefix (pokemon:133) before constructing a query
320 if ':' in name:
321 prefix_chunk, name = name.split(':', 1)
322 name = name.strip()
323
324 if not valid_types:
325 # Only use types from the query string if none were explicitly
326 # provided
327 prefixes = prefix_chunk.split(',')
328 valid_types = [_.strip() for _ in prefixes]
329
330 # Random lookup
331 if name == 'random':
332 return self.random_lookup(valid_types=valid_types)
333
334 # Do different things depending what the query looks like
335 # Note: Term objects do an exact match, so we don't have to worry about
336 # a query parser tripping on weird characters in the input
337 if '*' in name or '?' in name:
338 exact_only = True
339 query = whoosh.query.Wildcard(u'name', name)
340 elif rx_is_number.match(name):
341 # Don't spell-check numbers!
342 exact_only = True
343 query = whoosh.query.Term(u'row_id', name)
344 else:
345 # Not an integer
346 query = whoosh.query.Term(u'name', name)
347
348 ### Filter by type of object
349 type_terms = []
350 for valid_type in valid_types:
351 table_name = self._parse_table_name(valid_type)
352 if table_name:
353 # Quietly ignore bogus valid_types; more likely to DTRT
354 type_terms.append(whoosh.query.Term(u'table', table_name))
355
356 if type_terms:
357 query = query & whoosh.query.Or(type_terms)
358
359
360 ### Actual searching
361 searcher = self.index.searcher()
362 # XXX is this kosher? docs say search() takes a weighting arg, but it
363 # certainly does not
364 searcher.weighting = LanguageWeighting()
365 results = searcher.search(query,
366 limit=self.INTERMEDIATE_LOOKUP_RESULTS)
367
368 # Look for some fuzzy matches if necessary
369 if not exact_only and not results:
370 exact = False
371 results = []
372
373 for suggestion in self.speller.suggest(
374 name, self.INTERMEDIATE_LOOKUP_RESULTS):
375
376 query = whoosh.query.Term('name', suggestion)
377 results.extend(searcher.search(query))
378
379 ### Convert results to db objects
380 objects = self._whoosh_records_to_results(results, exact=exact)
381
382 # Only return up to 10 matches; beyond that, something is wrong. We
383 # strip out duplicate entries above, so it's remotely possible that we
384 # should have more than 10 here and lost a few. The speller returns 25
385 # to give us some padding, and should avoid that problem. Not a big
386 # deal if we lose the 25th-most-likely match anyway.
387 return objects[:self.MAX_LOOKUP_RESULTS]
388
389
390 def random_lookup(self, valid_types=[]):
391 """Returns a random lookup result from one of the provided
392 `valid_types`.
393 """
394
395 tables = []
396 for valid_type in valid_types:
397 table_name = self._parse_table_name(valid_type)
398 if table_name:
399 tables.append(self.indexed_tables[table_name])
400
401 if not tables:
402 # n.b.: It's possible we got a list of valid_types and none of them
403 # were valid, but this function is guaranteed to return
404 # *something*, so it politely selects from the entire index isntead
405 tables = self.indexed_tables.values()
406
407 # Rather than create an array of many hundred items and pick randomly
408 # from it, just pick a number up to the total number of potential
409 # items, then pick randomly from that, and partition the whole range
410 # into chunks. This also avoids the slight problem that the index
411 # contains more rows (for languages) for some items than others.
412 # XXX ought to cache this (in the index?) if possible
413 total = 0
414 partitions = []
415 for table in tables:
416 count = self.session.query(table).count()
417 total += count
418 partitions.append((table, count))
419
420 n = random.randint(1, total)
421 while n > partitions[0][1]:
422 n -= partitions[0][1]
423 partitions.pop(0)
424
425 return self.lookup(unicode(n), valid_types=[ partitions[0][0] ])
426
427 def prefix_lookup(self, prefix):
428 """Returns terms starting with the given exact prefix.
429
430 No special magic is currently done with the name; type prefixes are not
431 recognized.
432 """
433
434 query = whoosh.query.Prefix(u'name', self.normalize_name(prefix))
435
436 searcher = self.index.searcher()
437 searcher.weighting = LanguageWeighting()
438 results = searcher.search(query) # XXX , limit=self.MAX_LOOKUP_RESULTS)
439
440 return self._whoosh_records_to_results(results)