pokedex/lookup.py

   1 # encoding: utf8
   2 from collections import namedtuple
   3 import os, os.path
   4 import pkg_resources
   5 import re
   6
   7 from sqlalchemy.sql import func
   8 import whoosh
   9 import whoosh.filedb.filestore
  10 import whoosh.filedb.fileindex
  11 import whoosh.index
  12 from whoosh.qparser import QueryParser
  13 import whoosh.spelling
  14
  15 from pokedex.db import connect
  16 import pokedex.db.tables as tables
  17 from pokedex.roomaji import romanize
  18
  19 # Dictionary of table name => table class.
  20 # Need the table name so we can get the class from the table name after we
  21 # retrieve something from the index
  22 indexed_tables = {}
  23 for cls in [
  24         tables.Ability,
  25         tables.Item,
  26         tables.Move,
  27         tables.Pokemon,
  28         tables.Type,
  29     ]:
  30     indexed_tables[cls.__tablename__] = cls
  31
  32 # Dictionary of extra keys to file types of objects under, e.g. Pokémon can
  33 # also be looked up purely by number
  34 extra_keys = {
  35     tables.Move: [
  36         lambda row: u"move %d" % row.id,
  37     ],
  38     tables.Pokemon: [
  39         lambda row: unicode(row.id),
  40     ],
  41 }
  42
  43 def open_index(directory=None, session=None, recreate=False):
  44     """Opens the whoosh index stored in the named directory and returns (index,
  45     speller).  If the index doesn't already exist, it will be created.
  46
  47     `directory`
  48         Directory containing the index.  Defaults to a location within the
  49         `pokedex` egg directory.
  50
  51     `session`
  52         If the index needs to be created, this database session will be used.
  53         Defaults to an attempt to connect to the default SQLite database
  54         installed by `pokedex setup`.
  55
  56     `recreate`
  57         If set to True, the whoosh index will be created even if it already
  58         exists.
  59     """
  60
  61     # Defaults
  62     if not directory:
  63         directory = pkg_resources.resource_filename('pokedex',
  64                                                     'data/whoosh_index')
  65
  66     if not session:
  67         session = connect()
  68
  69     # Attempt to open or create the index
  70     directory_exists = os.path.exists(directory)
  71     if directory_exists and not recreate:
  72         # Already exists; should be an index!
  73         try:
  74             index = whoosh.index.open_dir(directory, indexname='MAIN')
  75             spell_store = whoosh.filedb.filestore.FileStorage(directory)
  76             speller = whoosh.spelling.SpellChecker(spell_store)
  77             return index, speller
  78         except whoosh.index.EmptyIndexError as e:
  79             # Apparently not a real index.  Fall out of the if and create it
  80             pass
  81
  82     if not directory_exists:
  83         os.mkdir(directory)
  84
  85
  86     # Create index
  87     schema = whoosh.fields.Schema(
  88         name=whoosh.fields.ID(stored=True),
  89         table=whoosh.fields.STORED,
  90         row_id=whoosh.fields.STORED,
  91         language=whoosh.fields.STORED,
  92     )
  93
  94     index = whoosh.index.create_in(directory, schema=schema, indexname='MAIN')
  95     writer = index.writer()
  96
  97     # Index every name in all our tables of interest
  98     # speller_entries becomes a list of (word, score) tuples; the score is 2
  99     # for English names, 1.5 for Roomaji, and 1 for everything else.  I think
 100     # this biases the results in the direction most people expect, especially
 101     # when e.g. German names are very similar to English names
 102     speller_entries = []
 103     for cls in indexed_tables.values():
 104         q = session.query(cls)
 105
 106         # Only index base Pokémon formes
 107         if hasattr(cls, 'forme_base_pokemon_id'):
 108             q = q.filter_by(forme_base_pokemon_id=None)
 109
 110         for row in q.yield_per(5):
 111             row_key = dict(table=cls.__tablename__, row_id=row.id)
 112
 113             name = row.name.lower()
 114             writer.add_document(name=name, **row_key)
 115             speller_entries.append((name, 1))
 116
 117             for extra_key_func in extra_keys.get(cls, []):
 118                 extra_key = extra_key_func(row)
 119                 writer.add_document(name=extra_key, **row_key)
 120
 121             # Pokemon also get other languages
 122             for foreign_name in getattr(row, 'foreign_names', []):
 123                 moonspeak = foreign_name.name.lower()
 124                 if name == moonspeak:
 125                     # Don't add the English name again as a different language;
 126                     # no point and it makes spell results confusing
 127                     continue
 128
 129                 writer.add_document(name=moonspeak,
 130                                     language=foreign_name.language.name,
 131                                     **row_key)
 132                 speller_entries.append((moonspeak, 3))
 133
 134                 # Add Roomaji too
 135                 if foreign_name.language.name == 'Japanese':
 136                     roomaji = romanize(foreign_name.name).lower()
 137                     writer.add_document(name=roomaji, language='Roomaji',
 138                                         **row_key)
 139                     speller_entries.append((roomaji, 8))
 140
 141
 142     writer.commit()
 143
 144     # Construct and populate a spell-checker index.  Quicker to do it all
 145     # at once, as every call to add_* does a commit(), and those seem to be
 146     # expensive
 147     speller = whoosh.spelling.SpellChecker(index.storage)
 148     speller.add_scored_words(speller_entries)
 149
 150     return index, speller
 151
 152
 153 LookupResult = namedtuple('LookupResult',
 154                           ['object', 'name', 'language', 'exact'])
 155 def lookup(name, session=None, indices=None, exact_only=False):
 156     """Attempts to find some sort of object, given a database session and name.
 157
 158     Returns a list of named (object, name, language, exact) tuples.  `object`
 159     is a database object, `name` is the name under which the object was found,
 160     `language` is the name of the language in which the name was found, and
 161     `exact` is True iff this was an exact match.
 162
 163     This function currently ONLY does fuzzy matching if there are no exact
 164     matches.
 165
 166     Formes are not returned; "Shaymin" will return only grass Shaymin.
 167
 168     Recognizes:
 169     - Pokémon names: "Eevee"
 170
 171     `name`
 172         Name of the thing to look for.
 173
 174     `session`
 175         A database session to use for retrieving objects.  As with get_index,
 176         if this is not provided, a connection to the default database will be
 177         attempted.
 178
 179     `indices`
 180         Tuple of index, speller as returned from `open_index()`.  Defaults to
 181         a call to `open_index()`.
 182
 183     `exact_only`
 184         If True, only exact matches are returned.  If set to False (the
 185         default), and the provided `name` doesn't match anything exactly,
 186         spelling correction will be attempted.
 187     """
 188
 189     if not session:
 190         session = connect()
 191
 192     if indices:
 193         index, speller = indices
 194     else:
 195         index, speller = open_index()
 196
 197     name = unicode(name)
 198
 199     exact = True
 200
 201     # Look for exact name.  A Term object does an exact match, so we don't have
 202     # to worry about a query parser tripping on weird characters in the input
 203     searcher = index.searcher()
 204     query = whoosh.query.Term('name', name.lower())
 205     results = searcher.search(query)
 206
 207     # Look for some fuzzy matches if necessary
 208     if not exact_only and not results:
 209         exact = False
 210         results = []
 211
 212         for suggestion in speller.suggest(name, 10):
 213             query = whoosh.query.Term('name', suggestion)
 214             results.extend(searcher.search(query))
 215
 216     ### Convert results to db objects
 217     objects = []
 218     seen = {}
 219     for result in results:
 220         # Skip dupe results
 221         # Note!  The speller prefers English names, but the query does not.  So
 222         # "latias" comes over "ratiasu".  "latias" matches only the English
 223         # row, comes out first, and all is well.
 224         # However!  The speller could then return "foo" which happens to be the
 225         # name for two different things in different languages, and the
 226         # non-English one could appear preferred.  This is not very likely.
 227         seen_key = result['table'], result['row_id']
 228         if seen_key in seen:
 229             continue
 230         seen[seen_key] = True
 231
 232         cls = indexed_tables[result['table']]
 233         obj = session.query(cls).get(result['row_id'])
 234         objects.append(LookupResult(object=obj,
 235                                     name=result['name'],
 236                                     language=result['language'],
 237                                     exact=exact))
 238
 239     return objects[:5]