2 from collections
import namedtuple
10 from sqlalchemy
.sql
import func
12 import whoosh
.filedb
.filestore
13 import whoosh
.filedb
.fileindex
15 from whoosh
.qparser
import QueryParser
17 import whoosh
.spelling
19 from pokedex
.db
import connect
20 import pokedex
.db
.tables
as tables
21 from pokedex
.roomaji
import romanize
23 __all__
= ['PokedexLookup']
26 rx_is_number
= re
.compile('^\d+$')
28 LookupResult
= namedtuple('LookupResult',
29 ['object', 'name', 'language', 'iso3166', 'exact'])
31 class LanguageWeighting(whoosh
.scoring
.Weighting
):
32 """A scoring class that forces otherwise-equal English results to come
33 before foreign results.
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
41 elif doc
['language'] == u
'Roomaji':
42 # Give Roomaji a little boost; it's most likely to be searched
45 # Everything else can drop down the totem pole
49 class PokedexLookup(object):
50 INTERMEDIATE_LOOKUP_RESULTS
= 25
51 MAX_LOOKUP_RESULTS
= 10
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
)
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.
73 Directory containing the index. Defaults to a location within the
74 `pokedex` egg directory.
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`.
82 If set to True, the whoosh index will be created even if it already
86 # By the time this returns, self.index, self.speller, and self.session
91 directory
= pkg_resources
.resource_filename('pokedex',
95 self
.session
= session
97 self
.session
= connect()
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.
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
)
108 except whoosh
.index
.EmptyIndexError
as e
:
109 # Apparently not a real index. Fall out and create it
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
120 if not directory_exists
:
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 forme_name
=whoosh
.fields
.ID
,
135 self
.index
= whoosh
.index
.create_in(directory
, schema
=schema
,
137 writer
= self
.index
.writer()
139 # Index every name in all our tables of interest
140 # speller_entries becomes a list of (word, score) tuples; the score is
141 # 2 for English names, 1.5 for Roomaji, and 1 for everything else. I
142 # think this biases the results in the direction most people expect,
143 # especially when e.g. German names are very similar to English names
145 for cls
in self
.indexed_tables
.values():
146 q
= session
.query(cls
)
148 for row
in q
.yield_per(5):
149 # Need to give forme_name a dummy value because I can't
150 # search for explicitly empty fields. Boo.
152 dict(table
=unicode(cls
.__tablename__
),
153 row_id
=unicode(row
.id),
154 forme_name
=u
'__empty__')
157 # If this is a form, mark it as such
158 # XXX foreign form names...?
159 if getattr(row
, 'forme_name', None):
160 # ...but if it's also the *default* form, index the name
162 if not getattr(row
, 'forme_base_pokemon_id', None):
163 new_key
= row_keys
[0].copy()
164 row_keys
.append(new_key
)
166 row_keys
[0]['forme_name'] = row
.forme_name
168 def add(name
, language
, iso3166
, score
):
169 normalized_name
= self
.normalize_name(name
)
170 for row_key
in row_keys
:
172 name
=normalized_name
, display_name
=name
,
173 language
=language
, iso3166
=iso3166
,
177 speller_entries
.append((normalized_name
, score
))
181 add(name
, None, u
'us', 1)
183 # Pokemon also get other languages
184 for foreign_name
in getattr(row
, 'foreign_names', []):
185 moonspeak
= foreign_name
.name
186 if name
== moonspeak
:
187 # Don't add the English name again as a different
188 # language; no point and it makes spell results
192 add(moonspeak
, foreign_name
.language
.name
,
193 foreign_name
.language
.iso3166
,
197 if foreign_name
.language
.name
== 'Japanese':
198 roomaji
= romanize(foreign_name
.name
)
199 add(roomaji
, u
'Roomaji', u
'jp', 8)
203 # Construct and populate a spell-checker index. Quicker to do it all
204 # at once, as every call to add_* does a commit(), and those seem to be
206 self
.speller
= whoosh
.spelling
.SpellChecker(self
.index
.storage
)
207 self
.speller
.add_scored_words(speller_entries
)
210 def normalize_name(self
, name
):
211 """Strips irrelevant formatting junk from name input.
213 Specifically: everything is lowercased, and accents are removed.
215 # http://stackoverflow.com/questions/517923/what-is-the-best-way-to-remove-accents-in-a-python-unicode-string
216 # Makes sense to me. Decompose by Unicode rules, then remove combining
217 # characters, then recombine. I'm explicitly doing it this way instead
218 # of testing combining() because Korean characters apparently
219 # decompose! But the results are considered letters, not combining
220 # characters, so testing for Mn works well, and combining them again
221 # makes them look right.
222 nkfd_form
= unicodedata
.normalize('NFKD', unicode(name
))
223 name
= u
"".join(c
for c
in nkfd_form
224 if unicodedata
.category(c
) != 'Mn')
225 name
= unicodedata
.normalize('NFC', name
)
233 def _parse_table_name(self
, name
):
234 """Takes a singular table name, table name, or table object and returns
237 Returns None for a bogus name.
239 if hasattr(name
, '__tablename__'):
240 return getattr(name
, '__tablename__')
241 elif name
in self
.indexed_tables
:
243 elif name
+ 's' in self
.indexed_tables
:
246 # Bogus. Be nice and return dummy
249 def _whoosh_records_to_results(self
, records
, exact
=True):
250 """Converts a list of whoosh's indexed records to LookupResult tuples
251 containing database objects.
253 # XXX this 'exact' thing is getting kinda leaky. would like a better
254 # way to handle it, since only lookup() cares about fuzzy results
257 for record
in records
:
259 seen_key
= record
['table'], record
['row_id']
262 seen
[seen_key
] = True
264 cls
= self
.indexed_tables
[record
['table']]
265 obj
= self
.session
.query(cls
).get(record
['row_id'])
267 results
.append(LookupResult(object=obj
,
268 name
=record
['display_name'],
269 language
=record
['language'],
270 iso3166
=record
['iso3166'],
276 def lookup(self
, input, valid_types
=[], exact_only
=False):
277 """Attempts to find some sort of object, given a name.
279 Returns a list of named (object, name, language, iso3166, exact)
280 tuples. `object` is a database object, `name` is the name under which
281 the object was found, `language` and `iso3166` are the name and country
282 code of the language in which the name was found, and `exact` is True
286 This function currently ONLY does fuzzy matching if there are no exact
289 Formes are not returned unless requested; "Shaymin" will return only
292 Extraneous whitespace is removed with extreme prejudice.
295 - Names: "Eevee", "Surf", "Run Away", "Payapa Berry", etc.
296 - Foreign names: "Iibui", "Eivui"
297 - Fuzzy names in whatever language: "Evee", "Ibui"
298 - IDs: "133", "192", "250"
300 - Type restrictions. "type:psychic" will only return the type. This
301 is how to make ID lookup useful. Multiple type specs can be entered
302 with commas, as "move,item:1". If `valid_types` are provided, any
303 type prefix will be ignored.
304 - Alternate formes can be specified merely like "wash rotom".
307 Name of the thing to look for.
310 A list of table objects or names, e.g., `['pokemon', 'moves']`. If
311 this is provided, only results in one of the given tables will be
315 If True, only exact matches are returned. If set to False (the
316 default), and the provided `name` doesn't match anything exactly,
317 spelling correction will be attempted.
320 name
= self
.normalize_name(input)
324 # Remove any type prefix (pokemon:133) before constructing a query
326 prefix_chunk
, name
= name
.split(':', 1)
330 # Only use types from the query string if none were explicitly
332 prefixes
= prefix_chunk
.split(',')
333 valid_types
= [_
.strip() for _
in prefixes
]
337 return self
.random_lookup(valid_types
=valid_types
)
339 # Do different things depending what the query looks like
340 # Note: Term objects do an exact match, so we don't have to worry about
341 # a query parser tripping on weird characters in the input
342 if '*' in name
or '?' in name
:
344 query
= whoosh
.query
.Wildcard(u
'name', name
)
345 elif rx_is_number
.match(name
):
346 # Don't spell-check numbers!
348 query
= whoosh
.query
.Term(u
'row_id', name
)
351 query
= whoosh
.query
.Term(u
'name', name
) \
352 & whoosh
.query
.Term(u
'forme_name', u
'__empty__')
354 # If there's a space in the input, this might be a form
356 form
, formless_name
= name
.split(' ', 1)
357 form_query
= whoosh
.query
.Term(u
'name', formless_name
) \
358 & whoosh
.query
.Term(u
'forme_name', form
)
359 query
= query | form_query
361 ### Filter by type of object
363 for valid_type
in valid_types
:
364 table_name
= self
._parse_table_name(valid_type
)
366 # Quietly ignore bogus valid_types; more likely to DTRT
367 type_terms
.append(whoosh
.query
.Term(u
'table', table_name
))
370 query
= query
& whoosh
.query
.Or(type_terms
)
374 searcher
= self
.index
.searcher()
375 # XXX is this kosher? docs say search() takes a weighting arg, but it
377 searcher
.weighting
= LanguageWeighting()
378 results
= searcher
.search(query
,
379 limit
=self
.INTERMEDIATE_LOOKUP_RESULTS
)
381 # Look for some fuzzy matches if necessary
382 if not exact_only
and not results
:
386 for suggestion
in self
.speller
.suggest(
387 name
, self
.INTERMEDIATE_LOOKUP_RESULTS
):
389 query
= whoosh
.query
.Term('name', suggestion
)
390 results
.extend(searcher
.search(query
))
392 ### Convert results to db objects
393 objects
= self
._whoosh_records_to_results(results
, exact
=exact
)
395 # Only return up to 10 matches; beyond that, something is wrong. We
396 # strip out duplicate entries above, so it's remotely possible that we
397 # should have more than 10 here and lost a few. The speller returns 25
398 # to give us some padding, and should avoid that problem. Not a big
399 # deal if we lose the 25th-most-likely match anyway.
400 return objects
[:self
.MAX_LOOKUP_RESULTS
]
403 def random_lookup(self
, valid_types
=[]):
404 """Returns a random lookup result from one of the provided
409 for valid_type
in valid_types
:
410 table_name
= self
._parse_table_name(valid_type
)
412 tables
.append(self
.indexed_tables
[table_name
])
415 # n.b.: It's possible we got a list of valid_types and none of them
416 # were valid, but this function is guaranteed to return
417 # *something*, so it politely selects from the entire index isntead
418 tables
= self
.indexed_tables
.values()
420 # Rather than create an array of many hundred items and pick randomly
421 # from it, just pick a number up to the total number of potential
422 # items, then pick randomly from that, and partition the whole range
423 # into chunks. This also avoids the slight problem that the index
424 # contains more rows (for languages) for some items than others.
425 # XXX ought to cache this (in the index?) if possible
429 count
= self
.session
.query(table
).count()
431 partitions
.append((table
, count
))
433 n
= random
.randint(1, total
)
434 while n
> partitions
[0][1]:
435 n
-= partitions
[0][1]
438 return self
.lookup(unicode(n
), valid_types
=[ partitions
[0][0] ])
440 def prefix_lookup(self
, prefix
):
441 """Returns terms starting with the given exact prefix.
443 No special magic is currently done with the name; type prefixes are not
447 query
= whoosh
.query
.Prefix(u
'name', self
.normalize_name(prefix
))
449 searcher
= self
.index
.searcher()
450 searcher
.weighting
= LanguageWeighting()
451 results
= searcher
.search(query
) # XXX , limit=self.MAX_LOOKUP_RESULTS)
453 return self
._whoosh_records_to_results(results
)