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.Location,
  62             tables.Move,
  63             tables.Nature,
  64             tables.Pokemon,
  65             tables.Type,
  66         )
  67     )
  68
  69
  70     def __init__(self, directory=None, session=None, recreate=False):
  71         """Opens the whoosh index stored in the named directory.  If the index
  72         doesn't already exist, it will be created.
  73
  74         `directory`
  75             Directory containing the index.  Defaults to a location within the
  76             `pokedex` egg directory.
  77
  78         `session`
  79             If the index needs to be created, this database session will be
  80             used.  Defaults to an attempt to connect to the default SQLite
  81             database installed by `pokedex setup`.
  82
  83         `recreate`
  84             If set to True, the whoosh index will be created even if it already
  85             exists.
  86         """
  87
  88         # By the time this returns, self.index, self.speller, and self.session
  89         # must be set
  90
  91         # Defaults
  92         if not directory:
  93             directory = pkg_resources.resource_filename('pokedex',
  94                                                         'data/whoosh-index')
  95
  96         if session:
  97             self.session = session
  98         else:
  99             self.session = connect()
 100
 101         # Attempt to open or create the index
 102         directory_exists = os.path.exists(directory)
 103         if directory_exists and not recreate:
 104             # Already exists; should be an index!  Bam, done.
 105             try:
 106                 self.index = whoosh.index.open_dir(directory, indexname='MAIN')
 107                 spell_store = whoosh.filedb.filestore.FileStorage(directory)
 108                 self.speller = whoosh.spelling.SpellChecker(spell_store)
 109                 return
 110             except whoosh.index.EmptyIndexError as e:
 111                 # Apparently not a real index.  Fall out and create it
 112                 pass
 113
 114         # Delete and start over if we're going to bail anyway.
 115         if directory_exists and recreate:
 116             # Be safe and only delete if it looks like a whoosh index, i.e.,
 117             # everything starts with _
 118             if all(f[0] == '_' for f in os.listdir(directory)):
 119                 shutil.rmtree(directory)
 120                 directory_exists = False
 121
 122         if not directory_exists:
 123             os.mkdir(directory)
 124
 125
 126         ### Create index
 127         schema = whoosh.fields.Schema(
 128             name=whoosh.fields.ID(stored=True),
 129             table=whoosh.fields.ID(stored=True),
 130             row_id=whoosh.fields.ID(stored=True),
 131             language=whoosh.fields.STORED,
 132             iso3166=whoosh.fields.STORED,
 133             display_name=whoosh.fields.STORED,  # non-lowercased name
 134         )
 135
 136         self.index = whoosh.index.create_in(directory, schema=schema,
 137                                             indexname='MAIN')
 138         writer = self.index.writer()
 139
 140         # Index every name in all our tables of interest
 141         # speller_entries becomes a list of (word, score) tuples; the score is
 142         # 2 for English names, 1.5 for Roomaji, and 1 for everything else.  I
 143         # think this biases the results in the direction most people expect,
 144         # especially when e.g. German names are very similar to English names
 145         speller_entries = []
 146         for cls in self.indexed_tables.values():
 147             q = self.session.query(cls)
 148
 149             for row in q.yield_per(5):
 150                 row_key = dict(table=unicode(cls.__tablename__),
 151                                row_id=unicode(row.id))
 152
 153                 def add(name, language, iso3166, score):
 154                     normalized_name = self.normalize_name(name)
 155
 156                     writer.add_document(
 157                         name=normalized_name, display_name=name,
 158                         language=language, iso3166=iso3166,
 159                         **row_key
 160                     )
 161
 162                     speller_entries.append((normalized_name, score))
 163
 164
 165                 # Add the basic English name to the index
 166                 if cls == tables.Pokemon:
 167                     # Pokémon need their form name added
 168                     # XXX kinda kludgy
 169                     add(row.full_name, None, u'us', 1)
 170
 171                     # If this is a default form, ALSO add the unadorned name,
 172                     # so 'Deoxys' alone will still do the right thing
 173                     if row.forme_name and not row.forme_base_pokemon_id:
 174                         add(row.name, None, u'us', 1)
 175                 else:
 176                     add(row.name, None, u'us', 1)
 177
 178                 # Some things also have other languages' names
 179                 # XXX other language form names..?
 180                 for foreign_name in getattr(row, 'foreign_names', []):
 181                     moonspeak = foreign_name.name
 182                     if row.name == moonspeak:
 183                         # Don't add the English name again as a different
 184                         # language; no point and it makes spell results
 185                         # confusing
 186                         continue
 187
 188                     add(moonspeak, foreign_name.language.name,
 189                                    foreign_name.language.iso3166,
 190                                    3)
 191
 192                     # Add Roomaji too
 193                     if foreign_name.language.name == 'Japanese':
 194                         roomaji = romanize(foreign_name.name)
 195                         add(roomaji, u'Roomaji', u'jp', 8)
 196
 197         writer.commit()
 198
 199         # Construct and populate a spell-checker index.  Quicker to do it all
 200         # at once, as every call to add_* does a commit(), and those seem to be
 201         # expensive
 202         self.speller = whoosh.spelling.SpellChecker(self.index.storage)
 203         self.speller.add_scored_words(speller_entries)
 204
 205
 206     def normalize_name(self, name):
 207         """Strips irrelevant formatting junk from name input.
 208
 209         Specifically: everything is lowercased, and accents are removed.
 210         """
 211         # http://stackoverflow.com/questions/517923/what-is-the-best-way-to-remove-accents-in-a-python-unicode-string
 212         # Makes sense to me.  Decompose by Unicode rules, then remove combining
 213         # characters, then recombine.  I'm explicitly doing it this way instead
 214         # of testing combining() because Korean characters apparently
 215         # decompose!  But the results are considered letters, not combining
 216         # characters, so testing for Mn works well, and combining them again
 217         # makes them look right.
 218         nkfd_form = unicodedata.normalize('NFKD', unicode(name))
 219         name = u"".join(c for c in nkfd_form
 220                         if unicodedata.category(c) != 'Mn')
 221         name = unicodedata.normalize('NFC', name)
 222
 223         name = name.strip()
 224         name = name.lower()
 225
 226         return name
 227
 228
 229     def _apply_valid_types(self, name, valid_types):
 230         """Combines the enforced `valid_types` with any from the search string
 231         itself and updates the query.
 232
 233         For example, a name of 'a,b:foo' and valid_types of b,c will search for
 234         only `b`s named "foo".
 235
 236         Returns `(name, merged_valid_types, term)`, where `name` has had any type
 237         prefix stripped, `merged_valid_types` combines the original
 238         `valid_types` with the type prefix, and `term` is a query term for
 239         limited to just the allowed types.  If there are no type restrictions
 240         at all, `term` will be None.
 241         """
 242
 243         # Remove any type prefix (pokemon:133) first
 244         user_valid_types = []
 245         if ':' in name:
 246             prefix_chunk, name = name.split(':', 1)
 247             name = name.strip()
 248
 249             prefixes = prefix_chunk.split(',')
 250             user_valid_types = [_.strip() for _ in prefixes]
 251
 252         # Merge the valid types together.  Only types that appear in BOTH lists
 253         # may be used.
 254         # As a special case, if the user asked for types that are explicitly
 255         # forbidden, completely ignore what the user requested
 256         combined_valid_types = []
 257         if user_valid_types and valid_types:
 258             combined_valid_types = list(
 259                 set(user_valid_types) & set(combined_valid_types)
 260             )
 261
 262             if not combined_valid_types:
 263                 # No overlap!  Just use the enforced ones
 264                 combined_valid_types = valid_types
 265         else:
 266             # One list or the other was blank, so just use the one that isn't
 267             combined_valid_types = valid_types + user_valid_types
 268
 269         if not combined_valid_types:
 270             # No restrictions
 271             return name, [], None
 272
 273         # Construct the term
 274         type_terms = []
 275         final_valid_types = []
 276         for valid_type in combined_valid_types:
 277             table_name = self._parse_table_name(valid_type)
 278
 279             # Quietly ignore bogus valid_types; more likely to DTRT
 280             if table_name:
 281                 final_valid_types.append(valid_type)
 282                 type_terms.append(whoosh.query.Term(u'table', table_name))
 283
 284         return name, final_valid_types, whoosh.query.Or(type_terms)
 285
 286
 287     def _parse_table_name(self, name):
 288         """Takes a singular table name, table name, or table object and returns
 289         the table name.
 290
 291         Returns None for a bogus name.
 292         """
 293         # Table object
 294         if hasattr(name, '__tablename__'):
 295             return getattr(name, '__tablename__')
 296
 297         # Table name
 298         for table in self.indexed_tables.values():
 299             if name in (table.__tablename__, table.__singlename__):
 300                 return table.__tablename__
 301
 302         # Bogus.  Be nice and return dummy
 303         return None
 304
 305     def _whoosh_records_to_results(self, records, exact=True):
 306         """Converts a list of whoosh's indexed records to LookupResult tuples
 307         containing database objects.
 308         """
 309         # XXX this 'exact' thing is getting kinda leaky.  would like a better
 310         # way to handle it, since only lookup() cares about fuzzy results
 311         seen = {}
 312         results = []
 313         for record in records:
 314             # Skip dupes
 315             seen_key = record['table'], record['row_id']
 316             if seen_key in seen:
 317                 continue
 318             seen[seen_key] = True
 319
 320             cls = self.indexed_tables[record['table']]
 321             obj = self.session.query(cls).get(record['row_id'])
 322
 323             results.append(LookupResult(object=obj,
 324                                         indexed_name=record['name'],
 325                                         name=record['display_name'],
 326                                         language=record['language'],
 327                                         iso3166=record['iso3166'],
 328                                         exact=exact))
 329
 330         return results
 331
 332
 333     def lookup(self, input, valid_types=[], exact_only=False):
 334         """Attempts to find some sort of object, given a name.
 335
 336         Returns a list of named (object, name, language, iso3166, exact)
 337         tuples.  `object` is a database object, `name` is the name under which
 338         the object was found, `language` and `iso3166` are the name and country
 339         code of the language in which the name was found, and `exact` is True
 340         iff this was an
 341         exact match.
 342
 343         This function currently ONLY does fuzzy matching if there are no exact
 344         matches.
 345
 346         Formes are not returned unless requested; "Shaymin" will return only
 347         grass Shaymin.
 348
 349         Extraneous whitespace is removed with extreme prejudice.
 350
 351         Recognizes:
 352         - Names: "Eevee", "Surf", "Run Away", "Payapa Berry", etc.
 353         - Foreign names: "Iibui", "Eivui"
 354         - Fuzzy names in whatever language: "Evee", "Ibui"
 355         - IDs: "133", "192", "250"
 356         Also:
 357         - Type restrictions.  "type:psychic" will only return the type.  This
 358           is how to make ID lookup useful.  Multiple type specs can be entered
 359           with commas, as "move,item:1".  If `valid_types` are provided, any
 360           type prefix will be ignored.
 361         - Alternate formes can be specified merely like "wash rotom".
 362
 363         `input`
 364             Name of the thing to look for.
 365
 366         `valid_types`
 367             A list of table objects or names, e.g., `['pokemon', 'moves']`.  If
 368             this is provided, only results in one of the given tables will be
 369             returned.
 370
 371         `exact_only`
 372             If True, only exact matches are returned.  If set to False (the
 373             default), and the provided `name` doesn't match anything exactly,
 374             spelling correction will be attempted.
 375         """
 376
 377         name = self.normalize_name(input)
 378         exact = True
 379         form = None
 380
 381         # Pop off any type prefix and merge with valid_types
 382         name, merged_valid_types, type_term = \
 383             self._apply_valid_types(name, valid_types)
 384
 385         # Random lookup
 386         if name == 'random':
 387             return self.random_lookup(valid_types=merged_valid_types)
 388
 389         # Do different things depending what the query looks like
 390         # Note: Term objects do an exact match, so we don't have to worry about
 391         # a query parser tripping on weird characters in the input
 392         try:
 393             # Let Python try to convert to a number, so 0xff works
 394             name_as_number = int(name, base=0)
 395         except ValueError:
 396             # Oh well
 397             name_as_number = None
 398
 399         if '*' in name or '?' in name:
 400             exact_only = True
 401             query = whoosh.query.Wildcard(u'name', name)
 402         elif name_as_number is not None:
 403             # Don't spell-check numbers!
 404             exact_only = True
 405             query = whoosh.query.Term(u'row_id', unicode(name_as_number))
 406         else:
 407             # Not an integer
 408             query = whoosh.query.Term(u'name', name)
 409
 410         if type_term:
 411             query = query & type_term
 412
 413
 414         ### Actual searching
 415         searcher = self.index.searcher()
 416         # XXX is this kosher?  docs say search() takes a weighting arg, but it
 417         # certainly does not
 418         searcher.weighting = LanguageWeighting()
 419         results = searcher.search(query,
 420                                   limit=self.INTERMEDIATE_LOOKUP_RESULTS)
 421
 422         # Look for some fuzzy matches if necessary
 423         if not exact_only and not results:
 424             exact = False
 425             results = []
 426
 427             for suggestion in self.speller.suggest(
 428                 name, self.INTERMEDIATE_LOOKUP_RESULTS):
 429
 430                 query = whoosh.query.Term('name', suggestion)
 431                 results.extend(searcher.search(query))
 432
 433         ### Convert results to db objects
 434         objects = self._whoosh_records_to_results(results, exact=exact)
 435
 436         # Only return up to 10 matches; beyond that, something is wrong.  We
 437         # strip out duplicate entries above, so it's remotely possible that we
 438         # should have more than 10 here and lost a few.  The speller returns 25
 439         # to give us some padding, and should avoid that problem.  Not a big
 440         # deal if we lose the 25th-most-likely match anyway.
 441         return objects[:self.MAX_LOOKUP_RESULTS]
 442
 443
 444     def random_lookup(self, valid_types=[]):
 445         """Returns a random lookup result from one of the provided
 446         `valid_types`.
 447         """
 448
 449         tables = []
 450         for valid_type in valid_types:
 451             table_name = self._parse_table_name(valid_type)
 452             if table_name:
 453                 tables.append(self.indexed_tables[table_name])
 454
 455         if not tables:
 456             # n.b.: It's possible we got a list of valid_types and none of them
 457             # were valid, but this function is guaranteed to return
 458             # *something*, so it politely selects from the entire index isntead
 459             tables = self.indexed_tables.values()
 460
 461         # Rather than create an array of many hundred items and pick randomly
 462         # from it, just pick a number up to the total number of potential
 463         # items, then pick randomly from that, and partition the whole range
 464         # into chunks.  This also avoids the slight problem that the index
 465         # contains more rows (for languages) for some items than others.
 466         # XXX ought to cache this (in the index?) if possible
 467         total = 0
 468         partitions = []
 469         for table in tables:
 470             count = self.session.query(table).count()
 471             total += count
 472             partitions.append((table, count))
 473
 474         n = random.randint(1, total)
 475         while n > partitions[0][1]:
 476             n -= partitions[0][1]
 477             partitions.pop(0)
 478
 479         return self.lookup(unicode(n), valid_types=[ partitions[0][0] ])
 480
 481     def prefix_lookup(self, prefix, valid_types=[]):
 482         """Returns terms starting with the given exact prefix.
 483
 484         Type prefixes are recognized, but no other name munging is done.
 485         """
 486
 487         # Pop off any type prefix and merge with valid_types
 488         prefix, merged_valid_types, type_term = \
 489             self._apply_valid_types(prefix, valid_types)
 490
 491         query = whoosh.query.Prefix(u'name', self.normalize_name(prefix))
 492
 493         if type_term:
 494             query = query & type_term
 495
 496         searcher = self.index.searcher()
 497         searcher.weighting = LanguageWeighting()
 498         results = searcher.search(query)  # XXX , limit=self.MAX_LOOKUP_RESULTS)
 499
 500         return self._whoosh_records_to_results(results)