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', 'indexed_name', '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
134 self
.index
= whoosh
.index
.create_in(directory
, schema
=schema
,
136 writer
= self
.index
.writer()
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
144 for cls
in self
.indexed_tables
.values():
145 q
= session
.query(cls
)
147 for row
in q
.yield_per(5):
148 row_key
= dict(table
=unicode(cls
.__tablename__
),
149 row_id
=unicode(row
.id))
151 def add(name
, language
, iso3166
, score
):
152 normalized_name
= self
.normalize_name(name
)
155 name
=normalized_name
, display_name
=name
,
156 language
=language
, iso3166
=iso3166
,
160 speller_entries
.append((normalized_name
, score
))
163 # Add the basic English name to the index
164 if cls
== tables
.Pokemon
:
165 # Pokémon need their form name added
167 add(row
.full_name
, None, u
'us', 1)
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)
174 add(row
.name
, None, u
'us', 1)
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
186 add(moonspeak
, foreign_name
.language
.name
,
187 foreign_name
.language
.iso3166
,
191 if foreign_name
.language
.name
== 'Japanese':
192 roomaji
= romanize(foreign_name
.name
)
193 add(roomaji
, u
'Roomaji', u
'jp', 8)
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
200 self
.speller
= whoosh
.spelling
.SpellChecker(self
.index
.storage
)
201 self
.speller
.add_scored_words(speller_entries
)
204 def normalize_name(self
, name
):
205 """Strips irrelevant formatting junk from name input.
207 Specifically: everything is lowercased, and accents are removed.
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
)
227 def _parse_table_name(self
, name
):
228 """Takes a singular table name, table name, or table object and returns
231 Returns None for a bogus name.
233 if hasattr(name
, '__tablename__'):
234 return getattr(name
, '__tablename__')
235 elif name
in self
.indexed_tables
:
237 elif name
+ 's' in self
.indexed_tables
:
240 # Bogus. Be nice and return dummy
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.
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
251 for record
in records
:
253 seen_key
= record
['table'], record
['row_id']
256 seen
[seen_key
] = True
258 cls
= self
.indexed_tables
[record
['table']]
259 obj
= self
.session
.query(cls
).get(record
['row_id'])
261 results
.append(LookupResult(object=obj
,
262 indexed_name
=record
['name'],
263 name
=record
['display_name'],
264 language
=record
['language'],
265 iso3166
=record
['iso3166'],
271 def lookup(self
, input, valid_types
=[], exact_only
=False):
272 """Attempts to find some sort of object, given a name.
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
281 This function currently ONLY does fuzzy matching if there are no exact
284 Formes are not returned unless requested; "Shaymin" will return only
287 Extraneous whitespace is removed with extreme prejudice.
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"
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".
302 Name of the thing to look for.
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
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.
315 name
= self
.normalize_name(input)
319 # Remove any type prefix (pokemon:133) before constructing a query
321 prefix_chunk
, name
= name
.split(':', 1)
325 # Only use types from the query string if none were explicitly
327 prefixes
= prefix_chunk
.split(',')
328 valid_types
= [_
.strip() for _
in prefixes
]
332 return self
.random_lookup(valid_types
=valid_types
)
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
:
339 query
= whoosh
.query
.Wildcard(u
'name', name
)
340 elif rx_is_number
.match(name
):
341 # Don't spell-check numbers!
343 query
= whoosh
.query
.Term(u
'row_id', name
)
346 query
= whoosh
.query
.Term(u
'name', name
)
348 ### Filter by type of object
350 for valid_type
in valid_types
:
351 table_name
= self
._parse_table_name(valid_type
)
353 # Quietly ignore bogus valid_types; more likely to DTRT
354 type_terms
.append(whoosh
.query
.Term(u
'table', table_name
))
357 query
= query
& whoosh
.query
.Or(type_terms
)
361 searcher
= self
.index
.searcher()
362 # XXX is this kosher? docs say search() takes a weighting arg, but it
364 searcher
.weighting
= LanguageWeighting()
365 results
= searcher
.search(query
,
366 limit
=self
.INTERMEDIATE_LOOKUP_RESULTS
)
368 # Look for some fuzzy matches if necessary
369 if not exact_only
and not results
:
373 for suggestion
in self
.speller
.suggest(
374 name
, self
.INTERMEDIATE_LOOKUP_RESULTS
):
376 query
= whoosh
.query
.Term('name', suggestion
)
377 results
.extend(searcher
.search(query
))
379 ### Convert results to db objects
380 objects
= self
._whoosh_records_to_results(results
, exact
=exact
)
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
]
390 def random_lookup(self
, valid_types
=[]):
391 """Returns a random lookup result from one of the provided
396 for valid_type
in valid_types
:
397 table_name
= self
._parse_table_name(valid_type
)
399 tables
.append(self
.indexed_tables
[table_name
])
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()
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
416 count
= self
.session
.query(table
).count()
418 partitions
.append((table
, count
))
420 n
= random
.randint(1, total
)
421 while n
> partitions
[0][1]:
422 n
-= partitions
[0][1]
425 return self
.lookup(unicode(n
), valid_types
=[ partitions
[0][0] ])
427 def prefix_lookup(self
, prefix
):
428 """Returns terms starting with the given exact prefix.
430 No special magic is currently done with the name; type prefixes are not
434 query
= whoosh
.query
.Prefix(u
'name', self
.normalize_name(prefix
))
436 searcher
= self
.index
.searcher()
437 searcher
.weighting
= LanguageWeighting()
438 results
= searcher
.search(query
) # XXX , limit=self.MAX_LOOKUP_RESULTS)
440 return self
._whoosh_records_to_results(results
)