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 = []
  99     for cls in indexed_tables.values():
 100         q = session.query(cls)
 101
 102         # Only index base Pokémon formes
 103         if hasattr(cls, 'forme_base_pokemon_id'):
 104             q = q.filter_by(forme_base_pokemon_id=None)
 105
 106         for row in q.yield_per(5):
 107             row_key = dict(table=cls.__tablename__, row_id=row.id)
 108
 109             name = row.name.lower()
 110             writer.add_document(name=name, **row_key)
 111             speller_entries.append(name)
 112
 113             for extra_key_func in extra_keys.get(cls, []):
 114                 extra_key = extra_key_func(row)
 115                 writer.add_document(name=extra_key, **row_key)
 116
 117             # Pokemon also get other languages
 118             if cls == tables.Pokemon:
 119                 for foreign_name in row.foreign_names:
 120                     name = foreign_name.name.lower()
 121                     writer.add_document(name=name,
 122                                         language=foreign_name.language.name,
 123                                         **row_key)
 124                     speller_entries.append(name)
 125
 126                     if foreign_name.language.name == 'Japanese':
 127                         # Add Roomaji too
 128                         roomaji = romanize(foreign_name.name).lower()
 129                         writer.add_document(name=roomaji,
 130                                             language='Roomaji',
 131                                             **row_key)
 132                         speller_entries.append(roomaji)
 133
 134
 135     writer.commit()
 136
 137     # Construct and populate a spell-checker index.  Quicker to do it all
 138     # at once, as every call to add_* does a commit(), and those seem to be
 139     # expensive
 140     speller = whoosh.spelling.SpellChecker(index.storage)
 141     speller.add_words(speller_entries)
 142
 143     return index, speller
 144
 145
 146 LookupResult = namedtuple('LookupResult', ['object', 'language', 'exact'])
 147 def lookup(name, session=None, indices=None, exact_only=False):
 148     """Attempts to find some sort of object, given a database session and name.
 149
 150     Returns a list of named (object, language, exact) tuples.  `object` is a
 151     database object, `language` is the name of the language in which the name
 152     was found, and `exact` is True iff this was an exact match.
 153
 154     This function currently ONLY does fuzzy matching if there are no exact
 155     matches.
 156
 157     Formes are not returned; "Shaymin" will return only grass Shaymin.
 158
 159     Recognizes:
 160     - Pokémon names: "Eevee"
 161
 162     `name`
 163         Name of the thing to look for.
 164
 165     `session`
 166         A database session to use for retrieving objects.  As with get_index,
 167         if this is not provided, a connection to the default database will be
 168         attempted.
 169
 170     `indices`
 171         Tuple of index, speller as returned from `open_index()`.  Defaults to
 172         a call to `open_index()`.
 173
 174     `exact_only`
 175         If True, only exact matches are returned.  If set to False (the
 176         default), and the provided `name` doesn't match anything exactly,
 177         spelling correction will be attempted.
 178     """
 179
 180     if not session:
 181         session = connect()
 182
 183     if indices:
 184         index, speller = indices
 185     else:
 186         index, speller = open_index()
 187
 188     name = unicode(name)
 189
 190     exact = True
 191
 192     # Look for exact name.  A Term object does an exact match, so we don't have
 193     # to worry about a query parser tripping on weird characters in the input
 194     searcher = index.searcher()
 195     query = whoosh.query.Term('name', name.lower())
 196     results = searcher.search(query)
 197
 198     # Look for some fuzzy matches if necessary
 199     if not exact_only and not results:
 200         exact = False
 201         results = []
 202
 203         for suggestion in speller.suggest(name, 10):
 204             query = whoosh.query.Term('name', suggestion)
 205             results.extend(searcher.search(query))
 206
 207     ### Convert results to db objects
 208     objects = []
 209     seen = {}
 210     for result in results:
 211         # Skip dupe results
 212         seen_key = result['table'], result['row_id']
 213         if seen_key in seen:
 214             continue
 215         seen[seen_key] = True
 216
 217         cls = indexed_tables[result['table']]
 218         obj = session.query(cls).get(result['row_id'])
 219         objects.append(LookupResult(obj, result['language'], exact))
 220
 221     return objects