veniasilente/veekun_pokedex
1
0
Fork 0
mirror of https://github.com/veekun/pokedex.git synced 2024-08-20 18:16:34 +00:00
Code Issues Projects Releases Wiki Activity

Totally overhauled lookup to use a class.

Browse source 
Now state is held within an object, rather than passed back to the
caller who must then pass it in again.  That was retarded and I don't
know why I ever did it.

Code is much cleaner now.

With apologies to anyone running annotate.
This commit is contained in:
Eevee 2010-03-21 23:27:47 -07:00
parent f9c17584f3
commit d6fd697018
2 changed files with 359 additions and 373 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
pokedex/__init__.py
View file

@ -54,11 +54,11 @@ def command_load(*args):
def command_setup(*args): def command_setup(*args):
session = connect() session = connect()
pokedex.db.load.load(session, verbose=False, drop_tables=True) pokedex.db.load.load(session, verbose=False, drop_tables=True)
pokedex.lookup.open_index(session=session, recreate=True) pokedex.lookup.PokedexLookup(session=session, recreate=True)
def command_lookup(name): def command_lookup(name):
results = pokedex.lookup.lookup(name) results = pokedex.lookup.PokedexLookup().lookup(name)
if not results: if not results:
print "No matches." print "No matches."
elif results[0].exact: elif results[0].exact:

330
pokedex/lookup.py
View file

@ -20,83 +20,93 @@ from pokedex.db import connect
import pokedex.db.tables as tables import pokedex.db.tables as tables
from pokedex.roomaji import romanize from pokedex.roomaji import romanize
__all__ = ['open_index', 'lookup', 'random_lookup'] __all__ = ['PokedexLookup']
INTERMEDIATE_LOOKUP_RESULTS = 25
MAX_LOOKUP_RESULTS = 10
# Dictionary of table name => table class. rx_is_number = re.compile('^\d+$')
# Need the table name so we can get the class from the table name after we
# retrieve something from the index LookupResult = namedtuple('LookupResult',
indexed_tables = {} ['object', 'name', 'language', 'iso3166', 'exact'])
for cls in [
class LanguageWeighting(whoosh.scoring.Weighting):
"""A scoring class that forces otherwise-equal English results to come
before foreign results.
"""
def score(self, searcher, fieldnum, text, docnum, weight, QTF=1):
doc = searcher.stored_fields(docnum)
if doc['language'] == None:
# English (well, "default"); leave it at 1
return weight
elif doc['language'] == u'Roomaji':
# Give Roomaji a little boost; it's most likely to be searched
return weight * 0.95
else:
# Everything else can drop down the totem pole
return weight * 0.9
class PokedexLookup(object):
INTERMEDIATE_LOOKUP_RESULTS = 25
MAX_LOOKUP_RESULTS = 10
# Dictionary of table name => table class.
# Need the table name so we can get the class from the table name after we
# retrieve something from the index
indexed_tables = dict(
(cls.__tablename__, cls)
for cls in (
tables.Ability, tables.Ability,
tables.Item, tables.Item,
tables.Move, tables.Move,
tables.Pokemon, tables.Pokemon,
tables.Type, tables.Type,
]: )
indexed_tables[cls.__tablename__] = cls )
def normalize(name):
"""Strips irrelevant formatting junk from name input.
Specifically: everything is lowercased, and accents are removed.
"""
# http://stackoverflow.com/questions/517923/what-is-the-best-way-to-remove-accents-in-a-python-unicode-string
# Makes sense to me. Decompose by Unicode rules, then remove combining
# characters, then recombine. I'm explicitly doing it this way instead of
# testing combining() because Korean characters apparently decompose! But
# the results are considered letters, not combining characters, so testing
# for Mn works well, and combining them again makes them look right.
nkfd_form = unicodedata.normalize('NFKD', unicode(name))
name = u"".join(c for c in nkfd_form
if unicodedata.category(c) != 'Mn')
name = unicodedata.normalize('NFC', name)
name = name.strip()
name = name.lower()
return name
def open_index(directory=None, session=None, recreate=False): def __init__(self, directory=None, session=None, recreate=False):
"""Opens the whoosh index stored in the named directory and returns (index, """Opens the whoosh index stored in the named directory. If the index
speller). If the index doesn't already exist, it will be created. doesn't already exist, it will be created.
`directory` `directory`
Directory containing the index. Defaults to a location within the Directory containing the index. Defaults to a location within the
`pokedex` egg directory. `pokedex` egg directory.
`session` `session`
If the index needs to be created, this database session will be used. If the index needs to be created, this database session will be
Defaults to an attempt to connect to the default SQLite database used. Defaults to an attempt to connect to the default SQLite
installed by `pokedex setup`. database installed by `pokedex setup`.
`recreate` `recreate`
If set to True, the whoosh index will be created even if it already If set to True, the whoosh index will be created even if it already
exists. exists.
""" """
# By the time this returns, self.index, self.speller, and self.session
# must be set
# Defaults # Defaults
if not directory: if not directory:
directory = pkg_resources.resource_filename('pokedex', directory = pkg_resources.resource_filename('pokedex',
'data/whoosh-index') 'data/whoosh-index')
if not session: if session:
session = connect() self.session = session
else:
self.session = connect()
# Attempt to open or create the index # Attempt to open or create the index
directory_exists = os.path.exists(directory) directory_exists = os.path.exists(directory)
if directory_exists and not recreate: if directory_exists and not recreate:
# Already exists; should be an index! # Already exists; should be an index! Bam, done.
try: try:
index = whoosh.index.open_dir(directory, indexname='MAIN') self.index = whoosh.index.open_dir(directory, indexname='MAIN')
spell_store = whoosh.filedb.filestore.FileStorage(directory) spell_store = whoosh.filedb.filestore.FileStorage(directory)
speller = whoosh.spelling.SpellChecker(spell_store) self.speller = whoosh.spelling.SpellChecker(spell_store)
return index, speller return
except whoosh.index.EmptyIndexError as e: except whoosh.index.EmptyIndexError as e:
# Apparently not a real index. Fall out of the if and create it # Apparently not a real index. Fall out and create it
pass pass
# Delete and start over if we're going to bail anyway. # Delete and start over if we're going to bail anyway.
@ -122,31 +132,33 @@ def open_index(directory=None, session=None, recreate=False):
forme_name=whoosh.fields.ID, forme_name=whoosh.fields.ID,
) )
index = whoosh.index.create_in(directory, schema=schema, indexname='MAIN') self.index = whoosh.index.create_in(directory, schema=schema,
writer = index.writer() indexname='MAIN')
writer = self.index.writer()
# Index every name in all our tables of interest # Index every name in all our tables of interest
# speller_entries becomes a list of (word, score) tuples; the score is 2 # speller_entries becomes a list of (word, score) tuples; the score is
# for English names, 1.5 for Roomaji, and 1 for everything else. I think # 2 for English names, 1.5 for Roomaji, and 1 for everything else. I
# this biases the results in the direction most people expect, especially # think this biases the results in the direction most people expect,
# when e.g. German names are very similar to English names # especially when e.g. German names are very similar to English names
speller_entries = [] speller_entries = []
for cls in indexed_tables.values(): for cls in self.indexed_tables.values():
q = session.query(cls) q = session.query(cls)
for row in q.yield_per(5): for row in q.yield_per(5):
# XXX need to give forme_name a dummy value because I can't search # Need to give forme_name a dummy value because I can't
# for explicitly empty fields. boo. # search for explicitly empty fields. Boo.
row_key = dict(table=unicode(cls.__tablename__), row_key = dict(table=unicode(cls.__tablename__),
row_id=unicode(row.id), row_id=unicode(row.id),
forme_name=u'XXX') forme_name=u'__empty__')
def add(name, language, iso3166, score): def add(name, language, iso3166, score):
normalized_name = normalize(name) normalized_name = self.normalize_name(name)
writer.add_document(name=normalized_name, display_name=name, writer.add_document(
language=language, name=normalized_name, display_name=name,
iso3166=iso3166, language=language, iso3166=iso3166,
**row_key) **row_key
)
speller_entries.append((normalized_name, score)) speller_entries.append((normalized_name, score))
# If this is a form, mark it as such # If this is a form, mark it as such
@ -160,8 +172,9 @@ def open_index(directory=None, session=None, recreate=False):
for foreign_name in getattr(row, 'foreign_names', []): for foreign_name in getattr(row, 'foreign_names', []):
moonspeak = foreign_name.name moonspeak = foreign_name.name
if name == moonspeak: if name == moonspeak:
# Don't add the English name again as a different language; # Don't add the English name again as a different
# no point and it makes spell results confusing # language; no point and it makes spell results
# confusing
continue continue
add(moonspeak, foreign_name.language.name, add(moonspeak, foreign_name.language.name,
@ -178,56 +191,55 @@ def open_index(directory=None, session=None, recreate=False):
# Construct and populate a spell-checker index. Quicker to do it all # Construct and populate a spell-checker index. Quicker to do it all
# at once, as every call to add_* does a commit(), and those seem to be # at once, as every call to add_* does a commit(), and those seem to be
# expensive # expensive
speller = whoosh.spelling.SpellChecker(index.storage) self.speller = whoosh.spelling.SpellChecker(self.index.storage)
speller.add_scored_words(speller_entries) self.speller.add_scored_words(speller_entries)
return index, speller
class LanguageWeighting(whoosh.scoring.Weighting): def normalize_name(self, name):
"""A scoring class that forces otherwise-equal English results to come """Strips irrelevant formatting junk from name input.
before foreign results.
Specifically: everything is lowercased, and accents are removed.
""" """
# http://stackoverflow.com/questions/517923/what-is-the-best-way-to-remove-accents-in-a-python-unicode-string
# Makes sense to me. Decompose by Unicode rules, then remove combining
# characters, then recombine. I'm explicitly doing it this way instead
# of testing combining() because Korean characters apparently
# decompose! But the results are considered letters, not combining
# characters, so testing for Mn works well, and combining them again
# makes them look right.
nkfd_form = unicodedata.normalize('NFKD', unicode(name))
name = u"".join(c for c in nkfd_form
if unicodedata.category(c) != 'Mn')
name = unicodedata.normalize('NFC', name)
def score(self, searcher, fieldnum, text, docnum, weight, QTF=1): name = name.strip()
doc = searcher.stored_fields(docnum) name = name.lower()
if doc['language'] == None:
# English (well, "default"); leave it at 1
return weight
elif doc['language'] == u'Roomaji':
# Give Roomaji a bit of a boost, as it's most likely to be searched
return weight * 0.95
else:
# Everything else can drop down the totem pole
return weight * 0.9
rx_is_number = re.compile('^\d+$') return name
LookupResult = namedtuple('LookupResult',
['object', 'name', 'language', 'iso3166', 'exact'])
def _parse_table_name(name): def _parse_table_name(self, name):
"""Takes a singular table name, table name, or table object and returns the """Takes a singular table name, table name, or table object and returns
table name. the table name.
Returns None for a bogus name. Returns None for a bogus name.
""" """
if hasattr(name, '__tablename__'): if hasattr(name, '__tablename__'):
return getattr(name, '__tablename__') return getattr(name, '__tablename__')
elif name in indexed_tables: elif name in self.indexed_tables:
return name return name
elif name + 's' in indexed_tables: elif name + 's' in self.indexed_tables:
return name + 's' return name + 's'
else: else:
# Bogus. Be nice and return dummy # Bogus. Be nice and return dummy
return None return None
def _whoosh_records_to_results(records, session, exact=True): def _whoosh_records_to_results(self, records, exact=True):
"""Converts a list of whoosh's indexed records to LookupResult tuples """Converts a list of whoosh's indexed records to LookupResult tuples
containing database objects. containing database objects.
""" """
# XXX this 'exact' thing is getting kinda leaky. would like a better way # XXX this 'exact' thing is getting kinda leaky. would like a better
# to handle it, since only lookup() cares about fuzzy results # way to handle it, since only lookup() cares about fuzzy results
seen = {} seen = {}
results = [] results = []
for record in records: for record in records:
@ -237,8 +249,8 @@ def _whoosh_records_to_results(records, session, exact=True):
continue continue
seen[seen_key] = True seen[seen_key] = True
cls = indexed_tables[record['table']] cls = self.indexed_tables[record['table']]
obj = session.query(cls).get(record['row_id']) obj = self.session.query(cls).get(record['row_id'])
results.append(LookupResult(object=obj, results.append(LookupResult(object=obj,
name=record['display_name'], name=record['display_name'],
@ -249,20 +261,21 @@ def _whoosh_records_to_results(records, session, exact=True):
return results return results
def lookup(input, valid_types=[], session=None, indices=None, exact_only=False): def lookup(self, input, valid_types=[], exact_only=False):
"""Attempts to find some sort of object, given a database session and name. """Attempts to find some sort of object, given a name.
Returns a list of named (object, name, language, iso3166, exact) tuples. Returns a list of named (object, name, language, iso3166, exact)
`object` is a database object, `name` is the name under which the object tuples. `object` is a database object, `name` is the name under which
was found, `language` and `iso3166` are the name and country code of the the object was found, `language` and `iso3166` are the name and country
language in which the name was found, and `exact` is True iff this was an code of the language in which the name was found, and `exact` is True
iff this was an
exact match. exact match.
This function currently ONLY does fuzzy matching if there are no exact This function currently ONLY does fuzzy matching if there are no exact
matches. matches.
Formes are not returned unless requested; "Shaymin" will return only grass Formes are not returned unless requested; "Shaymin" will return only
Shaymin. grass Shaymin.
Extraneous whitespace is removed with extreme prejudice. Extraneous whitespace is removed with extreme prejudice.
@ -272,10 +285,10 @@ def lookup(input, valid_types=[], session=None, indices=None, exact_only=False):
- Fuzzy names in whatever language: "Evee", "Ibui" - Fuzzy names in whatever language: "Evee", "Ibui"
- IDs: "133", "192", "250" - IDs: "133", "192", "250"
Also: Also:
- Type restrictions. "type:psychic" will only return the type. This is - Type restrictions. "type:psychic" will only return the type. This
how to make ID lookup useful. Multiple type specs can be entered with is how to make ID lookup useful. Multiple type specs can be entered
commas, as "move,item:1". If `valid_types` are provided, any type prefix with commas, as "move,item:1". If `valid_types` are provided, any
will be ignored. type prefix will be ignored.
- Alternate formes can be specified merely like "wash rotom". - Alternate formes can be specified merely like "wash rotom".
`input` `input`
@ -286,30 +299,13 @@ def lookup(input, valid_types=[], session=None, indices=None, exact_only=False):
this is provided, only results in one of the given tables will be this is provided, only results in one of the given tables will be
returned. returned.
`session`
A database session to use for retrieving objects. As with get_index,
if this is not provided, a connection to the default database will be
attempted.
`indices`
Tuple of index, speller as returned from `open_index()`. Defaults to
a call to `open_index()`.
`exact_only` `exact_only`
If True, only exact matches are returned. If set to False (the If True, only exact matches are returned. If set to False (the
default), and the provided `name` doesn't match anything exactly, default), and the provided `name` doesn't match anything exactly,
spelling correction will be attempted. spelling correction will be attempted.
""" """
if not session: name = self.normalize_name(input)
session = connect()
if indices:
index, speller = indices
else:
index, speller = open_index()
name = normalize(input)
exact = True exact = True
form = None form = None
@ -326,13 +322,11 @@ def lookup(input, valid_types=[], session=None, indices=None, exact_only=False):
# Random lookup # Random lookup
if name == 'random': if name == 'random':
return random_lookup(indices=(index, speller), return self.random_lookup(valid_types=valid_types)
session=session,
valid_types=valid_types)
# Do different things depending what the query looks like # Do different things depending what the query looks like
# Note: Term objects do an exact match, so we don't have to worry about a # Note: Term objects do an exact match, so we don't have to worry about
# query parser tripping on weird characters in the input # a query parser tripping on weird characters in the input
if '*' in name or '?' in name: if '*' in name or '?' in name:
exact_only = True exact_only = True
query = whoosh.query.Wildcard(u'name', name) query = whoosh.query.Wildcard(u'name', name)
@ -343,7 +337,7 @@ def lookup(input, valid_types=[], session=None, indices=None, exact_only=False):
else: else:
# Not an integer # Not an integer
query = whoosh.query.Term(u'name', name) \ query = whoosh.query.Term(u'name', name) \
& whoosh.query.Term(u'forme_name', u'XXX') & whoosh.query.Term(u'forme_name', u'__empty__')
# If there's a space in the input, this might be a form # If there's a space in the input, this might be a form
if ' ' in name: if ' ' in name:
@ -355,7 +349,7 @@ def lookup(input, valid_types=[], session=None, indices=None, exact_only=False):
### Filter by type of object ### Filter by type of object
type_terms = [] type_terms = []
for valid_type in valid_types: for valid_type in valid_types:
table_name = _parse_table_name(valid_type) table_name = self._parse_table_name(valid_type)
if table_name: if table_name:
# Quietly ignore bogus valid_types; more likely to DTRT # Quietly ignore bogus valid_types; more likely to DTRT
type_terms.append(whoosh.query.Term(u'table', table_name)) type_terms.append(whoosh.query.Term(u'table', table_name))
@ -365,59 +359,62 @@ def lookup(input, valid_types=[], session=None, indices=None, exact_only=False):
### Actual searching ### Actual searching
searcher = index.searcher() searcher = self.index.searcher()
searcher.weighting = LanguageWeighting() # XXX kosher? docs say search() # XXX is this kosher? docs say search() takes a weighting arg, but it
# takes a weighting kw but it
# certainly does not # certainly does not
results = searcher.search(query, limit=INTERMEDIATE_LOOKUP_RESULTS) searcher.weighting = LanguageWeighting()
results = searcher.search(query,
limit=self.INTERMEDIATE_LOOKUP_RESULTS)
# Look for some fuzzy matches if necessary # Look for some fuzzy matches if necessary
if not exact_only and not results: if not exact_only and not results:
exact = False exact = False
results = [] results = []
for suggestion in speller.suggest(name, INTERMEDIATE_LOOKUP_RESULTS): for suggestion in self.speller.suggest(
name, self.INTERMEDIATE_LOOKUP_RESULTS):
query = whoosh.query.Term('name', suggestion) query = whoosh.query.Term('name', suggestion)
results.extend(searcher.search(query)) results.extend(searcher.search(query))
### Convert results to db objects ### Convert results to db objects
objects = _whoosh_records_to_results(results, session, exact=exact) objects = self._whoosh_records_to_results(results, exact=exact)
# Only return up to 10 matches; beyond that, something is wrong. # Only return up to 10 matches; beyond that, something is wrong. We
# We strip out duplicate entries above, so it's remotely possible that we # strip out duplicate entries above, so it's remotely possible that we
# should have more than 10 here and lost a few. The speller returns 25 to # should have more than 10 here and lost a few. The speller returns 25
# give us some padding, and should avoid that problem. Not a big deal if # to give us some padding, and should avoid that problem. Not a big
# we lose the 25th-most-likely match anyway. # deal if we lose the 25th-most-likely match anyway.
return objects[:MAX_LOOKUP_RESULTS] return objects[:self.MAX_LOOKUP_RESULTS]
def random_lookup(valid_types=[], session=None, indices=None): def random_lookup(self, valid_types=[]):
"""Takes similar arguments as `lookup()`, but returns a random lookup """Returns a random lookup result from one of the provided
result from one of the provided `valid_types`. `valid_types`.
""" """
tables = [] tables = []
for valid_type in valid_types: for valid_type in valid_types:
table_name = _parse_table_name(valid_type) table_name = self._parse_table_name(valid_type)
if table_name: if table_name:
tables.append(indexed_tables[table_name]) tables.append(self.indexed_tables[table_name])
if not tables: if not tables:
# n.b.: It's possible we got a list of valid_types and none of them # n.b.: It's possible we got a list of valid_types and none of them
# were valid, but this function is guaranteed to return *something*, so # were valid, but this function is guaranteed to return
# it politely selects from the entire index isntead # *something*, so it politely selects from the entire index isntead
tables = indexed_tables.values() tables = self.indexed_tables.values()
# Rather than create an array of many hundred items and pick randomly from # Rather than create an array of many hundred items and pick randomly
# it, just pick a number up to the total number of potential items, then # from it, just pick a number up to the total number of potential
# pick randomly from that, and partition the whole range into chunks. # items, then pick randomly from that, and partition the whole range
# This also avoids the slight problem that the index contains more rows # into chunks. This also avoids the slight problem that the index
# (for languages) for some items than others. # contains more rows (for languages) for some items than others.
# XXX ought to cache this (in the index?) if possible # XXX ought to cache this (in the index?) if possible
total = 0 total = 0
partitions = [] partitions = []
for table in tables: for table in tables:
count = session.query(table).count() count = self.session.query(table).count()
total += count total += count
partitions.append((table, count)) partitions.append((table, count))
@ -426,30 +423,19 @@ def random_lookup(valid_types=[], session=None, indices=None):
n -= partitions[0][1] n -= partitions[0][1]
partitions.pop(0) partitions.pop(0)
return lookup(unicode(n), valid_types=[ partitions[0][0] ], return self.lookup(unicode(n), valid_types=[ partitions[0][0] ])
indices=indices, session=session)
def prefix_lookup(prefix, session=None, indices=None): def prefix_lookup(self, prefix):
"""Returns terms starting with the given exact prefix. """Returns terms starting with the given exact prefix.
No special magic is currently done with the name; type prefixes are not No special magic is currently done with the name; type prefixes are not
recognized. recognized.
`session` and `indices` are treated as with `lookup()`.
""" """
if not session: query = whoosh.query.Prefix(u'name', self.normalize_name(prefix))
session = connect()
if indices: searcher = self.index.searcher()
index, speller = indices
else:
index, speller = open_index()
query = whoosh.query.Prefix(u'name', normalize(prefix))
searcher = index.searcher()
searcher.weighting = LanguageWeighting() searcher.weighting = LanguageWeighting()
results = searcher.search(query) # XXX , limit=MAX_LOOKUP_RESULTS) results = searcher.search(query) # XXX , limit=self.MAX_LOOKUP_RESULTS)
return _whoosh_records_to_results(results, session) return self._whoosh_records_to_results(results)