2 from collections
import namedtuple
9 from sqlalchemy
.sql
import func
11 import whoosh
.filedb
.filestore
12 import whoosh
.filedb
.fileindex
14 from whoosh
.qparser
import QueryParser
16 import whoosh
.spelling
18 from pokedex
.db
import connect
19 import pokedex
.db
.tables
as tables
20 from pokedex
.roomaji
import romanize
22 __all__
= ['open_index', 'lookup', 'random_lookup']
24 INTERMEDIATE_LOOKUP_RESULTS
= 25
25 MAX_LOOKUP_RESULTS
= 10
27 # Dictionary of table name => table class.
28 # Need the table name so we can get the class from the table name after we
29 # retrieve something from the index
38 indexed_tables
[cls
.__tablename__
] = cls
40 def open_index(directory
=None, session
=None, recreate
=False):
41 """Opens the whoosh index stored in the named directory and returns (index,
42 speller). If the index doesn't already exist, it will be created.
45 Directory containing the index. Defaults to a location within the
46 `pokedex` egg directory.
49 If the index needs to be created, this database session will be used.
50 Defaults to an attempt to connect to the default SQLite database
51 installed by `pokedex setup`.
54 If set to True, the whoosh index will be created even if it already
60 directory
= pkg_resources
.resource_filename('pokedex',
66 # Attempt to open or create the index
67 directory_exists
= os
.path
.exists(directory
)
68 if directory_exists
and not recreate
:
69 # Already exists; should be an index!
71 index
= whoosh
.index
.open_dir(directory
, indexname
='MAIN')
72 spell_store
= whoosh
.filedb
.filestore
.FileStorage(directory
)
73 speller
= whoosh
.spelling
.SpellChecker(spell_store
)
75 except whoosh
.index
.EmptyIndexError
as e
:
76 # Apparently not a real index. Fall out of the if and create it
79 # Delete and start over if we're going to bail anyway.
80 if directory_exists
and recreate
:
81 # Be safe and only delete if it looks like a whoosh index, i.e.,
82 # everything starts with _
83 if all(f
[0] == '_' for f
in os
.listdir(directory
)):
84 shutil
.rmtree(directory
)
85 directory_exists
= False
87 if not directory_exists
:
92 schema
= whoosh
.fields
.Schema(
93 name
=whoosh
.fields
.ID(stored
=True),
94 table
=whoosh
.fields
.ID(stored
=True),
95 row_id
=whoosh
.fields
.ID(stored
=True),
96 language
=whoosh
.fields
.STORED
,
97 iso3166
=whoosh
.fields
.STORED
,
98 display_name
=whoosh
.fields
.STORED
, # non-lowercased name
99 forme_name
=whoosh
.fields
.ID
,
102 index
= whoosh
.index
.create_in(directory
, schema
=schema
, indexname
='MAIN')
103 writer
= index
.writer()
105 # Index every name in all our tables of interest
106 # speller_entries becomes a list of (word, score) tuples; the score is 2
107 # for English names, 1.5 for Roomaji, and 1 for everything else. I think
108 # this biases the results in the direction most people expect, especially
109 # when e.g. German names are very similar to English names
111 for cls
in indexed_tables
.values():
112 q
= session
.query(cls
)
114 for row
in q
.yield_per(5):
115 # XXX need to give forme_name a dummy value because I can't search
116 # for explicitly empty fields. boo.
117 row_key
= dict(table
=unicode(cls
.__tablename__
),
118 row_id
=unicode(row
.id),
121 def add(name
, language
, iso3166
, score
):
122 writer
.add_document(name
=name
.lower(), display_name
=name
,
126 speller_entries
.append((name
.lower(), score
))
128 # If this is a form, mark it as such
129 if getattr(row
, 'forme_base_pokemon_id', None):
130 row_key
['forme_name'] = row
.forme_name
133 add(name
, None, u
'us', 1)
135 # Pokemon also get other languages
136 for foreign_name
in getattr(row
, 'foreign_names', []):
137 moonspeak
= foreign_name
.name
138 if name
== moonspeak
:
139 # Don't add the English name again as a different language;
140 # no point and it makes spell results confusing
143 add(moonspeak
, foreign_name
.language
.name
,
144 foreign_name
.language
.iso3166
,
148 if foreign_name
.language
.name
== 'Japanese':
149 roomaji
= romanize(foreign_name
.name
)
150 add(roomaji
, u
'Roomaji', u
'jp', 8)
154 # Construct and populate a spell-checker index. Quicker to do it all
155 # at once, as every call to add_* does a commit(), and those seem to be
157 speller
= whoosh
.spelling
.SpellChecker(index
.storage
)
158 speller
.add_scored_words(speller_entries
)
160 return index
, speller
163 class LanguageWeighting(whoosh
.scoring
.Weighting
):
164 """A scoring class that forces otherwise-equal English results to come
165 before foreign results.
168 def score(self
, searcher
, fieldnum
, text
, docnum
, weight
, QTF
=1):
169 doc
= searcher
.stored_fields(docnum
)
170 if doc
['language'] == None:
171 # English (well, "default"); leave it at 1
173 elif doc
['language'] == u
'Roomaji':
174 # Give Roomaji a bit of a boost, as it's most likely to be searched
177 # Everything else can drop down the totem pole
180 rx_is_number
= re
.compile('^\d+$')
182 LookupResult
= namedtuple('LookupResult',
183 ['object', 'name', 'language', 'iso3166', 'exact'])
185 def _parse_table_name(name
):
186 """Takes a singular table name, table name, or table object and returns the
189 Returns None for a bogus name.
191 if hasattr(name
, '__tablename__'):
192 return getattr(name
, '__tablename__')
193 elif name
in indexed_tables
:
195 elif name
+ 's' in indexed_tables
:
198 # Bogus. Be nice and return dummy
201 def _whoosh_records_to_results(records
, session
, exact
=True):
202 """Converts a list of whoosh's indexed records to LookupResult tuples
203 containing database objects.
205 # XXX this 'exact' thing is getting kinda leaky. would like a better way
206 # to handle it, since only lookup() cares about fuzzy results
209 for record
in records
:
211 seen_key
= record
['table'], record
['row_id']
214 seen
[seen_key
] = True
216 cls
= indexed_tables
[record
['table']]
217 obj
= session
.query(cls
).get(record
['row_id'])
219 results
.append(LookupResult(object=obj
,
220 name
=record
['display_name'],
221 language
=record
['language'],
222 iso3166
=record
['iso3166'],
228 def lookup(input, valid_types
=[], session
=None, indices
=None, exact_only
=False):
229 """Attempts to find some sort of object, given a database session and name.
231 Returns a list of named (object, name, language, iso3166, exact) tuples.
232 `object` is a database object, `name` is the name under which the object
233 was found, `language` and `iso3166` are the name and country code of the
234 language in which the name was found, and `exact` is True iff this was an
237 This function currently ONLY does fuzzy matching if there are no exact
240 Formes are not returned unless requested; "Shaymin" will return only grass
243 Extraneous whitespace is removed with extreme prejudice.
246 - Names: "Eevee", "Surf", "Run Away", "Payapa Berry", etc.
247 - Foreign names: "Iibui", "Eivui"
248 - Fuzzy names in whatever language: "Evee", "Ibui"
249 - IDs: "133", "192", "250"
251 - Type restrictions. "type:psychic" will only return the type. This is
252 how to make ID lookup useful. Multiple type specs can be entered with
253 commas, as "move,item:1". If `valid_types` are provided, any type prefix
255 - Alternate formes can be specified merely like "wash rotom".
258 Name of the thing to look for.
261 A list of table objects or names, e.g., `['pokemon', 'moves']`. If
262 this is provided, only results in one of the given tables will be
266 A database session to use for retrieving objects. As with get_index,
267 if this is not provided, a connection to the default database will be
271 Tuple of index, speller as returned from `open_index()`. Defaults to
272 a call to `open_index()`.
275 If True, only exact matches are returned. If set to False (the
276 default), and the provided `name` doesn't match anything exactly,
277 spelling correction will be attempted.
284 index
, speller
= indices
286 index
, speller
= open_index()
288 name
= unicode(input).strip().lower()
292 # Remove any type prefix (pokemon:133) before constructing a query
294 prefix_chunk
, name
= name
.split(':', 1)
298 # Only use types from the query string if none were explicitly
300 prefixes
= prefix_chunk
.split(',')
301 valid_types
= [_
.strip() for _
in prefixes
]
305 return random_lookup(indices
=(index
, speller
),
307 valid_types
=valid_types
)
309 # Do different things depending what the query looks like
310 # Note: Term objects do an exact match, so we don't have to worry about a
311 # query parser tripping on weird characters in the input
312 if '*' in name
or '?' in name
:
314 query
= whoosh
.query
.Wildcard(u
'name', name
)
315 elif rx_is_number
.match(name
):
316 # Don't spell-check numbers!
318 query
= whoosh
.query
.Term(u
'row_id', name
)
321 query
= whoosh
.query
.Term(u
'name', name
) \
322 & whoosh
.query
.Term(u
'forme_name', u
'XXX')
324 # If there's a space in the input, this might be a form
326 form
, formless_name
= name
.split(' ', 1)
327 form_query
= whoosh
.query
.Term(u
'name', formless_name
) \
328 & whoosh
.query
.Term(u
'forme_name', form
)
329 query
= query | form_query
331 ### Filter by type of object
333 for valid_type
in valid_types
:
334 table_name
= _parse_table_name(valid_type
)
336 # Quietly ignore bogus valid_types; more likely to DTRT
337 type_terms
.append(whoosh
.query
.Term(u
'table', table_name
))
340 query
= query
& whoosh
.query
.Or(type_terms
)
344 searcher
= index
.searcher()
345 searcher
.weighting
= LanguageWeighting() # XXX kosher? docs say search()
346 # takes a weighting kw but it
348 results
= searcher
.search(query
, limit
=INTERMEDIATE_LOOKUP_RESULTS
)
350 # Look for some fuzzy matches if necessary
351 if not exact_only
and not results
:
355 for suggestion
in speller
.suggest(name
, INTERMEDIATE_LOOKUP_RESULTS
):
356 query
= whoosh
.query
.Term('name', suggestion
)
357 results
.extend(searcher
.search(query
))
359 ### Convert results to db objects
360 objects
= _whoosh_records_to_results(results
, session
, exact
=exact
)
362 # Only return up to 10 matches; beyond that, something is wrong.
363 # We strip out duplicate entries above, so it's remotely possible that we
364 # should have more than 10 here and lost a few. The speller returns 25 to
365 # give us some padding, and should avoid that problem. Not a big deal if
366 # we lose the 25th-most-likely match anyway.
367 return objects
[:MAX_LOOKUP_RESULTS
]
370 def random_lookup(valid_types
=[], session
=None, indices
=None):
371 """Takes similar arguments as `lookup()`, but returns a random lookup
372 result from one of the provided `valid_types`.
376 for valid_type
in valid_types
:
377 table_name
= _parse_table_name(valid_type
)
379 tables
.append(indexed_tables
[table_name
])
382 # n.b.: It's possible we got a list of valid_types and none of them
383 # were valid, but this function is guaranteed to return *something*, so
384 # it politely selects from the entire index isntead
385 tables
= indexed_tables
.values()
387 # Rather than create an array of many hundred items and pick randomly from
388 # it, just pick a number up to the total number of potential items, then
389 # pick randomly from that, and partition the whole range into chunks.
390 # This also avoids the slight problem that the index contains more rows
391 # (for languages) for some items than others.
392 # XXX ought to cache this (in the index?) if possible
396 count
= session
.query(table
).count()
398 partitions
.append((table
, count
))
400 n
= random
.randint(1, total
)
401 while n
> partitions
[0][1]:
402 n
-= partitions
[0][1]
405 return lookup(unicode(n
), valid_types
=[ partitions
[0][0] ],
406 indices
=indices
, session
=session
)
408 def prefix_lookup(prefix
, session
=None, indices
=None):
409 """Returns terms starting with the given exact prefix.
411 No special magic is currently done with the name; type prefixes are not
414 `session` and `indices` are treated as with `lookup()`.
421 index
, speller
= indices
423 index
, speller
= open_index()
425 query
= whoosh
.query
.Prefix(u
'name', prefix
.lower())
427 searcher
= index
.searcher()
428 searcher
.weighting
= LanguageWeighting()
429 results
= searcher
.search(query
) # XXX , limit=MAX_LOOKUP_RESULTS)
431 return _whoosh_records_to_results(results
, session
)