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__
= ['open_index', 'lookup', 'random_lookup']
25 INTERMEDIATE_LOOKUP_RESULTS
= 25
26 MAX_LOOKUP_RESULTS
= 10
28 # Dictionary of table name => table class.
29 # Need the table name so we can get the class from the table name after we
30 # retrieve something from the index
39 indexed_tables
[cls
.__tablename__
] = cls
42 """Strips irrelevant formatting junk from name input.
44 Specifically: everything is lowercased, and accents are removed.
46 # http://stackoverflow.com/questions/517923/what-is-the-best-way-to-remove-accents-in-a-python-unicode-string
47 # Makes sense to me. Decompose by Unicode rules, then remove combining
48 # characters, then recombine. I'm explicitly doing it this way instead of
49 # testing combining() because Korean characters apparently decompose! But
50 # the results are considered letters, not combining characters, so testing
51 # for Mn works well, and combining them again makes them look right.
52 nkfd_form
= unicodedata
.normalize('NFKD', unicode(name
))
53 name
= u
"".join(c
for c
in nkfd_form
54 if unicodedata
.category(c
) != 'Mn')
55 name
= unicodedata
.normalize('NFC', name
)
63 def open_index(directory
=None, session
=None, recreate
=False):
64 """Opens the whoosh index stored in the named directory and returns (index,
65 speller). If the index doesn't already exist, it will be created.
68 Directory containing the index. Defaults to a location within the
69 `pokedex` egg directory.
72 If the index needs to be created, this database session will be used.
73 Defaults to an attempt to connect to the default SQLite database
74 installed by `pokedex setup`.
77 If set to True, the whoosh index will be created even if it already
83 directory
= pkg_resources
.resource_filename('pokedex',
89 # Attempt to open or create the index
90 directory_exists
= os
.path
.exists(directory
)
91 if directory_exists
and not recreate
:
92 # Already exists; should be an index!
94 index
= whoosh
.index
.open_dir(directory
, indexname
='MAIN')
95 spell_store
= whoosh
.filedb
.filestore
.FileStorage(directory
)
96 speller
= whoosh
.spelling
.SpellChecker(spell_store
)
98 except whoosh
.index
.EmptyIndexError
as e
:
99 # Apparently not a real index. Fall out of the if and create it
102 # Delete and start over if we're going to bail anyway.
103 if directory_exists
and recreate
:
104 # Be safe and only delete if it looks like a whoosh index, i.e.,
105 # everything starts with _
106 if all(f
[0] == '_' for f
in os
.listdir(directory
)):
107 shutil
.rmtree(directory
)
108 directory_exists
= False
110 if not directory_exists
:
115 schema
= whoosh
.fields
.Schema(
116 name
=whoosh
.fields
.ID(stored
=True),
117 table
=whoosh
.fields
.ID(stored
=True),
118 row_id
=whoosh
.fields
.ID(stored
=True),
119 language
=whoosh
.fields
.STORED
,
120 iso3166
=whoosh
.fields
.STORED
,
121 display_name
=whoosh
.fields
.STORED
, # non-lowercased name
122 forme_name
=whoosh
.fields
.ID
,
125 index
= whoosh
.index
.create_in(directory
, schema
=schema
, indexname
='MAIN')
126 writer
= index
.writer()
128 # Index every name in all our tables of interest
129 # speller_entries becomes a list of (word, score) tuples; the score is 2
130 # for English names, 1.5 for Roomaji, and 1 for everything else. I think
131 # this biases the results in the direction most people expect, especially
132 # when e.g. German names are very similar to English names
134 for cls
in indexed_tables
.values():
135 q
= session
.query(cls
)
137 for row
in q
.yield_per(5):
138 # XXX need to give forme_name a dummy value because I can't search
139 # for explicitly empty fields. boo.
140 row_key
= dict(table
=unicode(cls
.__tablename__
),
141 row_id
=unicode(row
.id),
144 def add(name
, language
, iso3166
, score
):
145 normalized_name
= normalize(name
)
146 writer
.add_document(name
=normalized_name
, display_name
=name
,
150 speller_entries
.append((normalized_name
, score
))
152 # If this is a form, mark it as such
153 if getattr(row
, 'forme_base_pokemon_id', None):
154 row_key
['forme_name'] = row
.forme_name
157 add(name
, None, u
'us', 1)
159 # Pokemon also get other languages
160 for foreign_name
in getattr(row
, 'foreign_names', []):
161 moonspeak
= foreign_name
.name
162 if name
== moonspeak
:
163 # Don't add the English name again as a different language;
164 # no point and it makes spell results confusing
167 add(moonspeak
, foreign_name
.language
.name
,
168 foreign_name
.language
.iso3166
,
172 if foreign_name
.language
.name
== 'Japanese':
173 roomaji
= romanize(foreign_name
.name
)
174 add(roomaji
, u
'Roomaji', u
'jp', 8)
178 # Construct and populate a spell-checker index. Quicker to do it all
179 # at once, as every call to add_* does a commit(), and those seem to be
181 speller
= whoosh
.spelling
.SpellChecker(index
.storage
)
182 speller
.add_scored_words(speller_entries
)
184 return index
, speller
187 class LanguageWeighting(whoosh
.scoring
.Weighting
):
188 """A scoring class that forces otherwise-equal English results to come
189 before foreign results.
192 def score(self
, searcher
, fieldnum
, text
, docnum
, weight
, QTF
=1):
193 doc
= searcher
.stored_fields(docnum
)
194 if doc
['language'] == None:
195 # English (well, "default"); leave it at 1
197 elif doc
['language'] == u
'Roomaji':
198 # Give Roomaji a bit of a boost, as it's most likely to be searched
201 # Everything else can drop down the totem pole
204 rx_is_number
= re
.compile('^\d+$')
206 LookupResult
= namedtuple('LookupResult',
207 ['object', 'name', 'language', 'iso3166', 'exact'])
209 def _parse_table_name(name
):
210 """Takes a singular table name, table name, or table object and returns the
213 Returns None for a bogus name.
215 if hasattr(name
, '__tablename__'):
216 return getattr(name
, '__tablename__')
217 elif name
in indexed_tables
:
219 elif name
+ 's' in indexed_tables
:
222 # Bogus. Be nice and return dummy
225 def _whoosh_records_to_results(records
, session
, exact
=True):
226 """Converts a list of whoosh's indexed records to LookupResult tuples
227 containing database objects.
229 # XXX this 'exact' thing is getting kinda leaky. would like a better way
230 # to handle it, since only lookup() cares about fuzzy results
233 for record
in records
:
235 seen_key
= record
['table'], record
['row_id']
238 seen
[seen_key
] = True
240 cls
= indexed_tables
[record
['table']]
241 obj
= session
.query(cls
).get(record
['row_id'])
243 results
.append(LookupResult(object=obj
,
244 name
=record
['display_name'],
245 language
=record
['language'],
246 iso3166
=record
['iso3166'],
252 def lookup(input, valid_types
=[], session
=None, indices
=None, exact_only
=False):
253 """Attempts to find some sort of object, given a database session and name.
255 Returns a list of named (object, name, language, iso3166, exact) tuples.
256 `object` is a database object, `name` is the name under which the object
257 was found, `language` and `iso3166` are the name and country code of the
258 language in which the name was found, and `exact` is True iff this was an
261 This function currently ONLY does fuzzy matching if there are no exact
264 Formes are not returned unless requested; "Shaymin" will return only grass
267 Extraneous whitespace is removed with extreme prejudice.
270 - Names: "Eevee", "Surf", "Run Away", "Payapa Berry", etc.
271 - Foreign names: "Iibui", "Eivui"
272 - Fuzzy names in whatever language: "Evee", "Ibui"
273 - IDs: "133", "192", "250"
275 - Type restrictions. "type:psychic" will only return the type. This is
276 how to make ID lookup useful. Multiple type specs can be entered with
277 commas, as "move,item:1". If `valid_types` are provided, any type prefix
279 - Alternate formes can be specified merely like "wash rotom".
282 Name of the thing to look for.
285 A list of table objects or names, e.g., `['pokemon', 'moves']`. If
286 this is provided, only results in one of the given tables will be
290 A database session to use for retrieving objects. As with get_index,
291 if this is not provided, a connection to the default database will be
295 Tuple of index, speller as returned from `open_index()`. Defaults to
296 a call to `open_index()`.
299 If True, only exact matches are returned. If set to False (the
300 default), and the provided `name` doesn't match anything exactly,
301 spelling correction will be attempted.
308 index
, speller
= indices
310 index
, speller
= open_index()
312 name
= normalize(input)
316 # Remove any type prefix (pokemon:133) before constructing a query
318 prefix_chunk
, name
= name
.split(':', 1)
322 # Only use types from the query string if none were explicitly
324 prefixes
= prefix_chunk
.split(',')
325 valid_types
= [_
.strip() for _
in prefixes
]
329 return random_lookup(indices
=(index
, speller
),
331 valid_types
=valid_types
)
333 # Do different things depending what the query looks like
334 # Note: Term objects do an exact match, so we don't have to worry about a
335 # query parser tripping on weird characters in the input
336 if '*' in name
or '?' in name
:
338 query
= whoosh
.query
.Wildcard(u
'name', name
)
339 elif rx_is_number
.match(name
):
340 # Don't spell-check numbers!
342 query
= whoosh
.query
.Term(u
'row_id', name
)
345 query
= whoosh
.query
.Term(u
'name', name
) \
346 & whoosh
.query
.Term(u
'forme_name', u
'XXX')
348 # If there's a space in the input, this might be a form
350 form
, formless_name
= name
.split(' ', 1)
351 form_query
= whoosh
.query
.Term(u
'name', formless_name
) \
352 & whoosh
.query
.Term(u
'forme_name', form
)
353 query
= query | form_query
355 ### Filter by type of object
357 for valid_type
in valid_types
:
358 table_name
= _parse_table_name(valid_type
)
360 # Quietly ignore bogus valid_types; more likely to DTRT
361 type_terms
.append(whoosh
.query
.Term(u
'table', table_name
))
364 query
= query
& whoosh
.query
.Or(type_terms
)
368 searcher
= index
.searcher()
369 searcher
.weighting
= LanguageWeighting() # XXX kosher? docs say search()
370 # takes a weighting kw but it
372 results
= searcher
.search(query
, limit
=INTERMEDIATE_LOOKUP_RESULTS
)
374 # Look for some fuzzy matches if necessary
375 if not exact_only
and not results
:
379 for suggestion
in speller
.suggest(name
, INTERMEDIATE_LOOKUP_RESULTS
):
380 query
= whoosh
.query
.Term('name', suggestion
)
381 results
.extend(searcher
.search(query
))
383 ### Convert results to db objects
384 objects
= _whoosh_records_to_results(results
, session
, exact
=exact
)
386 # Only return up to 10 matches; beyond that, something is wrong.
387 # We strip out duplicate entries above, so it's remotely possible that we
388 # should have more than 10 here and lost a few. The speller returns 25 to
389 # give us some padding, and should avoid that problem. Not a big deal if
390 # we lose the 25th-most-likely match anyway.
391 return objects
[:MAX_LOOKUP_RESULTS
]
394 def random_lookup(valid_types
=[], session
=None, indices
=None):
395 """Takes similar arguments as `lookup()`, but returns a random lookup
396 result from one of the provided `valid_types`.
400 for valid_type
in valid_types
:
401 table_name
= _parse_table_name(valid_type
)
403 tables
.append(indexed_tables
[table_name
])
406 # n.b.: It's possible we got a list of valid_types and none of them
407 # were valid, but this function is guaranteed to return *something*, so
408 # it politely selects from the entire index isntead
409 tables
= indexed_tables
.values()
411 # Rather than create an array of many hundred items and pick randomly from
412 # it, just pick a number up to the total number of potential items, then
413 # pick randomly from that, and partition the whole range into chunks.
414 # This also avoids the slight problem that the index contains more rows
415 # (for languages) for some items than others.
416 # XXX ought to cache this (in the index?) if possible
420 count
= session
.query(table
).count()
422 partitions
.append((table
, count
))
424 n
= random
.randint(1, total
)
425 while n
> partitions
[0][1]:
426 n
-= partitions
[0][1]
429 return lookup(unicode(n
), valid_types
=[ partitions
[0][0] ],
430 indices
=indices
, session
=session
)
432 def prefix_lookup(prefix
, session
=None, indices
=None):
433 """Returns terms starting with the given exact prefix.
435 No special magic is currently done with the name; type prefixes are not
438 `session` and `indices` are treated as with `lookup()`.
445 index
, speller
= indices
447 index
, speller
= open_index()
449 query
= whoosh
.query
.Prefix(u
'name', normalize(prefix
))
451 searcher
= index
.searcher()
452 searcher
.weighting
= LanguageWeighting()
453 results
= searcher
.search(query
) # XXX , limit=MAX_LOOKUP_RESULTS)
455 return _whoosh_records_to_results(results
, session
)