pokedex/lookup.py

   1 # encoding: utf8
   2 from collections import namedtuple
   3 import os, os.path
   4 import pkg_resources
   5 import random
   6 import re
   7 import shutil
   8 import unicodedata
   9
  10 from sqlalchemy.sql import func
  11 import whoosh
  12 import whoosh.filedb.filestore
  13 import whoosh.filedb.fileindex
  14 import whoosh.index
  15 from whoosh.qparser import QueryParser
  16 import whoosh.scoring
  17 import whoosh.spelling
  18
  19 from pokedex.db import connect
  20 import pokedex.db.tables as tables
  21 from pokedex.roomaji import romanize
  22
  23 __all__ = ['PokedexLookup']
  24
  25
  26 rx_is_number = re.compile('^\d+$')
  27
  28 LookupResult = namedtuple('LookupResult',
  29     ['object', 'indexed_name', 'name', 'language', 'iso3166', 'exact'])
  30
  31 class LanguageWeighting(whoosh.scoring.Weighting):
  32     """A scoring class that forces otherwise-equal English results to come
  33     before foreign results.
  34     """
  35
  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
  40             return weight
  41         elif doc['language'] == u'Roomaji':
  42             # Give Roomaji a little boost; it's most likely to be searched
  43             return weight * 0.95
  44         else:
  45             # Everything else can drop down the totem pole
  46             return weight * 0.9
  47
  48
  49 class PokedexLookup(object):
  50     INTERMEDIATE_LOOKUP_RESULTS = 25
  51     MAX_LOOKUP_RESULTS = 10
  52
  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)
  58         for cls in (
  59             tables.Ability,
  60             tables.Item,
  61             tables.Move,
  62             tables.Pokemon,
  63             tables.Type,
  64         )
  65     )
  66
  67
  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.
  71
  72         `directory`
  73             Directory containing the index.  Defaults to a location within the
  74             `pokedex` egg directory.
  75
  76         `session`
  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`.
  80
  81         `recreate`
  82             If set to True, the whoosh index will be created even if it already
  83             exists.
  84         """
  85
  86         # By the time this returns, self.index, self.speller, and self.session
  87         # must be set
  88
  89         # Defaults
  90         if not directory:
  91             directory = pkg_resources.resource_filename('pokedex',
  92                                                         'data/whoosh-index')
  93
  94         if session:
  95             self.session = session
  96         else:
  97             self.session = connect()
  98
  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.
 103             try:
 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)
 107                 return
 108             except whoosh.index.EmptyIndexError as e:
 109                 # Apparently not a real index.  Fall out and create it
 110                 pass
 111
 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
 119
 120         if not directory_exists:
 121             os.mkdir(directory)
 122
 123
 124         ### Create index
 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         )
 133
 134         self.index = whoosh.index.create_in(directory, schema=schema,
 135                                             indexname='MAIN')
 136         writer = self.index.writer()
 137
 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
 143         speller_entries = []
 144         for cls in self.indexed_tables.values():
 145             q = session.query(cls)
 146
 147             for row in q.yield_per(5):
 148                 row_key = dict(table=unicode(cls.__tablename__),
 149                                row_id=unicode(row.id))
 150
 151                 def add(name, language, iso3166, score):
 152                     normalized_name = self.normalize_name(name)
 153
 154                     writer.add_document(
 155                         name=normalized_name, display_name=name,
 156                         language=language, iso3166=iso3166,
 157                         **row_key
 158                     )
 159
 160                     speller_entries.append((normalized_name, score))
 161
 162
 163                 # Add the basic English name to the index
 164                 if cls == tables.Pokemon:
 165                     # Pokémon need their form name added
 166                     # XXX kinda kludgy
 167                     add(row.full_name, None, u'us', 1)
 168
 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)
 173                 else:
 174                     add(row.name, None, u'us', 1)
 175
 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
 183                         # confusing
 184                         continue
 185
 186                     add(moonspeak, foreign_name.language.name,
 187                                    foreign_name.language.iso3166,
 188                                    3)
 189
 190                     # Add Roomaji too
 191                     if foreign_name.language.name == 'Japanese':
 192                         roomaji = romanize(foreign_name.name)
 193                         add(roomaji, u'Roomaji', u'jp', 8)
 194
 195         writer.commit()
 196
 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
 199         # expensive
 200         self.speller = whoosh.spelling.SpellChecker(self.index.storage)
 201         self.speller.add_scored_words(speller_entries)
 202
 203
 204     def normalize_name(self, name):
 205         """Strips irrelevant formatting junk from name input.
 206
 207         Specifically: everything is lowercased, and accents are removed.
 208         """
 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)
 220
 221         name = name.strip()
 222         name = name.lower()
 223
 224         return name
 225
 226
 227     def _parse_table_name(self, name):
 228         """Takes a singular table name, table name, or table object and returns
 229         the table name.
 230
 231         Returns None for a bogus name.
 232         """
 233         if hasattr(name, '__tablename__'):
 234             return getattr(name, '__tablename__')
 235         elif name in self.indexed_tables:
 236             return name
 237         elif name + 's' in self.indexed_tables:
 238             return name + 's'
 239         else:
 240             # Bogus.  Be nice and return dummy
 241             return None
 242
 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.
 246         """
 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
 249         seen = {}
 250         results = []
 251         for record in records:
 252             # Skip dupes
 253             seen_key = record['table'], record['row_id']
 254             if seen_key in seen:
 255                 continue
 256             seen[seen_key] = True
 257
 258             cls = self.indexed_tables[record['table']]
 259             obj = self.session.query(cls).get(record['row_id'])
 260
 261             results.append(LookupResult(object=obj,
 262                                         indexed_name=record['name'],
 263                                         name=record['display_name'],
 264                                         language=record['language'],
 265                                         iso3166=record['iso3166'],
 266                                         exact=exact))
 267
 268         return results
 269
 270
 271     def lookup(self, input, valid_types=[], exact_only=False):
 272         """Attempts to find some sort of object, given a name.
 273
 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
 278         iff this was an
 279         exact match.
 280
 281         This function currently ONLY does fuzzy matching if there are no exact
 282         matches.
 283
 284         Formes are not returned unless requested; "Shaymin" will return only
 285         grass Shaymin.
 286
 287         Extraneous whitespace is removed with extreme prejudice.
 288
 289         Recognizes:
 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"
 294         Also:
 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".
 300
 301         `input`
 302             Name of the thing to look for.
 303
 304         `valid_types`
 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
 307             returned.
 308
 309         `exact_only`
 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.
 313         """
 314
 315         name = self.normalize_name(input)
 316         exact = True
 317         form = None
 318
 319         # Remove any type prefix (pokemon:133) before constructing a query
 320         if ':' in name:
 321             prefix_chunk, name = name.split(':', 1)
 322             name = name.strip()
 323
 324             if not valid_types:
 325                 # Only use types from the query string if none were explicitly
 326                 # provided
 327                 prefixes = prefix_chunk.split(',')
 328                 valid_types = [_.strip() for _ in prefixes]
 329
 330         # Random lookup
 331         if name == 'random':
 332             return self.random_lookup(valid_types=valid_types)
 333
 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:
 338             exact_only = True
 339             query = whoosh.query.Wildcard(u'name', name)
 340         elif rx_is_number.match(name):
 341             # Don't spell-check numbers!
 342             exact_only = True
 343             query = whoosh.query.Term(u'row_id', name)
 344         else:
 345             # Not an integer
 346             query = whoosh.query.Term(u'name', name) \
 347                   & whoosh.query.Term(u'forme_name', u'__empty__')
 348
 349             # If there's a space in the input, this might be a form
 350             if ' ' in name:
 351                 form, formless_name = name.split(' ', 1)
 352                 form_query = whoosh.query.Term(u'name', formless_name) \
 353                            & whoosh.query.Term(u'forme_name', form)
 354                 query = query | form_query
 355
 356         ### Filter by type of object
 357         type_terms = []
 358         for valid_type in valid_types:
 359             table_name = self._parse_table_name(valid_type)
 360             if table_name:
 361                 # Quietly ignore bogus valid_types; more likely to DTRT
 362                 type_terms.append(whoosh.query.Term(u'table', table_name))
 363
 364         if type_terms:
 365             query = query & whoosh.query.Or(type_terms)
 366
 367
 368         ### Actual searching
 369         searcher = self.index.searcher()
 370         # XXX is this kosher?  docs say search() takes a weighting arg, but it
 371         # certainly does not
 372         searcher.weighting = LanguageWeighting()
 373         results = searcher.search(query,
 374                                   limit=self.INTERMEDIATE_LOOKUP_RESULTS)
 375
 376         # Look for some fuzzy matches if necessary
 377         if not exact_only and not results:
 378             exact = False
 379             results = []
 380
 381             for suggestion in self.speller.suggest(
 382                 name, self.INTERMEDIATE_LOOKUP_RESULTS):
 383
 384                 query = whoosh.query.Term('name', suggestion)
 385                 results.extend(searcher.search(query))
 386
 387         ### Convert results to db objects
 388         objects = self._whoosh_records_to_results(results, exact=exact)
 389
 390         # Only return up to 10 matches; beyond that, something is wrong.  We
 391         # strip out duplicate entries above, so it's remotely possible that we
 392         # should have more than 10 here and lost a few.  The speller returns 25
 393         # to give us some padding, and should avoid that problem.  Not a big
 394         # deal if we lose the 25th-most-likely match anyway.
 395         return objects[:self.MAX_LOOKUP_RESULTS]
 396
 397
 398     def random_lookup(self, valid_types=[]):
 399         """Returns a random lookup result from one of the provided
 400         `valid_types`.
 401         """
 402
 403         tables = []
 404         for valid_type in valid_types:
 405             table_name = self._parse_table_name(valid_type)
 406             if table_name:
 407                 tables.append(self.indexed_tables[table_name])
 408
 409         if not tables:
 410             # n.b.: It's possible we got a list of valid_types and none of them
 411             # were valid, but this function is guaranteed to return
 412             # *something*, so it politely selects from the entire index isntead
 413             tables = self.indexed_tables.values()
 414
 415         # Rather than create an array of many hundred items and pick randomly
 416         # from it, just pick a number up to the total number of potential
 417         # items, then pick randomly from that, and partition the whole range
 418         # into chunks.  This also avoids the slight problem that the index
 419         # contains more rows (for languages) for some items than others.
 420         # XXX ought to cache this (in the index?) if possible
 421         total = 0
 422         partitions = []
 423         for table in tables:
 424             count = self.session.query(table).count()
 425             total += count
 426             partitions.append((table, count))
 427
 428         n = random.randint(1, total)
 429         while n > partitions[0][1]:
 430             n -= partitions[0][1]
 431             partitions.pop(0)
 432
 433         return self.lookup(unicode(n), valid_types=[ partitions[0][0] ])
 434
 435     def prefix_lookup(self, prefix):
 436         """Returns terms starting with the given exact prefix.
 437
 438         No special magic is currently done with the name; type prefixes are not
 439         recognized.
 440         """
 441
 442         query = whoosh.query.Prefix(u'name', self.normalize_name(prefix))
 443
 444         searcher = self.index.searcher()
 445         searcher.weighting = LanguageWeighting()
 446         results = searcher.search(query)  # XXX , limit=self.MAX_LOOKUP_RESULTS)
 447
 448         return self._whoosh_records_to_results(results)