X-Git-Url: http://git.veekun.com/zzz-pokedex.git/blobdiff_plain/0b37ca809fe5f1c97b76f4e3e8a398caa40935fb..b7268b5c97e6c0f5c2b92e5befb53d1a251e42b3:/pokedex/lookup.py?ds=sidebyside diff --git a/pokedex/lookup.py b/pokedex/lookup.py index b50871b..55e1b5a 100644 --- a/pokedex/lookup.py +++ b/pokedex/lookup.py @@ -1,11 +1,17 @@ # encoding: utf8 +import os, os.path +import pkg_resources import re from sqlalchemy.sql import func import whoosh +import whoosh.filedb.filestore +import whoosh.filedb.fileindex +import whoosh.index from whoosh.qparser import QueryParser import whoosh.spelling +from pokedex.db import connect import pokedex.db.tables as tables # Dictionary of table name => table class. @@ -13,38 +19,81 @@ import pokedex.db.tables as tables # retrieve something from the index indexed_tables = {} for cls in [ + tables.Ability, + tables.Item, + tables.Move, tables.Pokemon, + tables.Type, ]: indexed_tables[cls.__tablename__] = cls # Dictionary of extra keys to file types of objects under, e.g. Pokémon can # also be looked up purely by number extra_keys = { + tables.Move: [ + lambda row: u"move %d" % row.id, + ], tables.Pokemon: [ lambda row: unicode(row.id), ], } -index_bits = {} -def get_index(session): - """Returns (index, speller). +def open_index(directory=None, session=None, recreate=False): + """Opens the whoosh index stored in the named directory and returns (index, + speller). If the index doesn't already exist, it will be created. + + `directory` + Directory containing the index. Defaults to a location within the + `pokedex` egg directory. - Creates an index if one does not exist. + `session` + If the index needs to be created, this database session will be used. + Defaults to an attempt to connect to the default SQLite database + installed by `pokedex setup`. + + `recreate` + If set to True, the whoosh index will be created even if it already + exists. """ - if index_bits: - return index_bits['index'], index_bits['speller'] + # Defaults + if not directory: + directory = pkg_resources.resource_filename('pokedex', + 'data/whoosh_index') + + if not session: + session = connect() - store = whoosh.store.RamStorage() + # Attempt to open or create the index + directory_exists = os.path.exists(directory) + if directory_exists and not recreate: + # Already exists; should be an index! + try: + index = whoosh.index.open_dir(directory, indexname='pokedex') + speller = whoosh.index.open_dir(directory, indexname='spelling') + return index, speller + except whoosh.index.EmptyIndexError as e: + # Apparently not a real index. Fall out of the if and create it + pass + + if not directory_exists: + os.mkdir(directory) + + + # Create index schema = whoosh.fields.Schema( name=whoosh.fields.ID(stored=True), table=whoosh.fields.STORED, row_id=whoosh.fields.STORED, - language_id=whoosh.fields.STORED, + language=whoosh.fields.STORED, + + # Whoosh 0.2 explodes when using a file-stored schema with no TEXT + # columns. Appease it + dummy=whoosh.fields.TEXT, ) - # Construct a straight lookup index - index = whoosh.index.Index(store, schema=schema, create=True) + index = whoosh.index.create_in(directory, schema=schema, + indexname='pokedex') writer = index.writer() # Index every name in all our tables of interest @@ -69,16 +118,24 @@ def get_index(session): speller_entries.append(name) - for extra_key_func in extra_keys[cls]: + for extra_key_func in extra_keys.get(cls, []): extra_key = extra_key_func(row) writer.add_document(name=extra_key, **row_key) writer.commit() + # XXX GIHWEGREHKG + old__schema = whoosh.spelling.SpellChecker._schema + def new__schema(self): + schema = old__schema(self) + schema.add('dummy', whoosh.fields.TEXT) + return schema + whoosh.spelling.SpellChecker._schema = new__schema + # 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 # expensive - speller = whoosh.spelling.SpellChecker(index.storage) + speller = whoosh.spelling.SpellChecker(index.storage, indexname='spelling') # WARNING: HERE BE DRAGONS # whoosh.spelling refuses to index things that don't look like words. # Unfortunately, this doesn't work so well for Pokémon (Mr. Mime, @@ -86,7 +143,7 @@ def get_index(session): # complications. # The below is copied from SpellChecker.add_scored_words without the check # for isalpha(). XXX get whoosh patched to make this unnecessary! - writer = whoosh.writing.IndexWriter(speller.index()) + writer = speller.index(create=True).writer() for word in speller_entries: fields = {"word": word, "score": 1} for size in xrange(speller.mingram, speller.maxgram + 1): @@ -100,12 +157,10 @@ def get_index(session): writer.commit() # end copy-pasta - index_bits['index'] = index - index_bits['speller'] = speller - index_bits['store'] = store - return index_bits['index'], index_bits['speller'] + return index, speller + -def lookup(session, name, exact_only=False): +def lookup(name, session=None, exact_only=False): """Attempts to find some sort of object, given a database session and name. Returns (objects, exact) where `objects` is a list of database objects, and @@ -117,16 +172,32 @@ def lookup(session, name, exact_only=False): Currently recognizes: - Pokémon names: "Eevee" + + `name` + Name of the thing to look for. + + `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. + + `exact_only` + If True, only exact matches are returned. If set to False (the + default), and the provided `name` doesn't match anything exactly, + spelling correction will be attempted. """ - exact = True + if not session: + session = connect() - index, speller = get_index(session) + index, speller = open_index() + + exact = True # Look for exact name. A Term object does an exact match, so we don't have # to worry about a query parser tripping on weird characters in the input searcher = index.searcher() - query = whoosh.query.Term('name', name) + query = whoosh.query.Term('name', name.lower()) results = searcher.search(query) if not exact_only: