Fix Arceus's mismatched and missing sprites. #189

[zzz-pokedex.git] / pokedex / lookup.py

diff --git a/pokedex/lookup.py b/pokedex/lookup.py
index ce5bb76..8c3e444 100644 (file)
--- a/pokedex/lookup.py
+++ b/pokedex/lookup.py
@@ -1,7 +1,5 @@
 # encoding: utf8
 # encoding: utf8

-from collections import namedtuple

 import os, os.path
 import os, os.path

-import pkg_resources

 import random
 import re
 import shutil
 import random
 import re
 import shutil


@@ -16,9 +14,12 @@ from whoosh.qparser import QueryParser
 import whoosh.scoring
 import whoosh.spelling
 
 import whoosh.scoring
 import whoosh.spelling
 

+from pokedex.util import namedtuple
+

 from pokedex.db import connect
 import pokedex.db.tables as tables
 from pokedex.roomaji import romanize
 from pokedex.db import connect
 import pokedex.db.tables as tables
 from pokedex.roomaji import romanize

+from pokedex.defaults import get_default_index_dir

 
 __all__ = ['PokedexLookup']
 
 
 __all__ = ['PokedexLookup']
 


@@ -28,28 +29,66 @@ rx_is_number = re.compile('^\d+$')
 LookupResult = namedtuple('LookupResult',
     ['object', 'indexed_name', 'name', 'language', 'iso3166', 'exact'])
 
 LookupResult = namedtuple('LookupResult',
     ['object', 'indexed_name', 'name', 'language', 'iso3166', 'exact'])
 

+class UninitializedIndex(object):
+    class UninitializedIndexError(Exception):
+        pass
+
+    def __nonzero__(self):
+        """Dummy object should identify itself as False."""
+        return False
+
+    def __bool__(self):
+        """Python 3000 version of the above.  Future-proofing rules!"""
+        return False
+
+    def __getattr__(self, *args, **kwargs):
+        raise self.UninitializedIndexError(
+            "The lookup index does not exist.  Please use `pokedex setup` "
+            "or lookup.rebuild_index() to create it."
+        )
+

 class LanguageWeighting(whoosh.scoring.Weighting):
     """A scoring class that forces otherwise-equal English results to come
     before foreign results.
     """
 
 class LanguageWeighting(whoosh.scoring.Weighting):
     """A scoring class that forces otherwise-equal English results to come
     before foreign results.
     """
 

+    def __init__(self, extra_weights={}, *args, **kwargs):
+        """`extra_weights` may be a dictionary of weights which will be
+        factored in.
+
+        Intended for use with spelling corrections, which come along with their
+        own weightings.
+        """
+        self.extra_weights = extra_weights
+        super(LanguageWeighting, self).__init__(*args, **kwargs)
+

     def score(self, searcher, fieldnum, text, docnum, weight, QTF=1):
         doc = searcher.stored_fields(docnum)
     def score(self, searcher, fieldnum, text, docnum, weight, QTF=1):
         doc = searcher.stored_fields(docnum)

+
+        # Apply extra weight
+        weight = weight * self.extra_weights.get(text, 1.0)
+

         if doc['language'] == None:
             # English (well, "default"); leave it at 1
             return weight
         elif doc['language'] == u'Roomaji':
             # Give Roomaji a little boost; it's most likely to be searched
         if doc['language'] == None:
             # English (well, "default"); leave it at 1
             return weight
         elif doc['language'] == u'Roomaji':
             # Give Roomaji a little boost; it's most likely to be searched

-            return weight * 0.95
+            return weight * 0.9

         else:
             # Everything else can drop down the totem pole
         else:
             # Everything else can drop down the totem pole

-            return weight * 0.9
+            return weight * 0.8

 
 
 class PokedexLookup(object):
     INTERMEDIATE_LOOKUP_RESULTS = 25
     MAX_LOOKUP_RESULTS = 10
 
 
 
 class PokedexLookup(object):
     INTERMEDIATE_LOOKUP_RESULTS = 25
     MAX_LOOKUP_RESULTS = 10
 

+    # The speller only checks how much the input matches a word; there can be
+    # all manner of extra unmatched junk, and it won't affect the weighting.
+    # To compensate, greatly boost the weighting of matches at the beginning
+    # and end, so nearly-full-word-matches are much better
+    SPELLER_OPTIONS = dict(booststart=10.0, boostend=9.0)
+

     # Dictionary of table name => table class.
     # Need the table name so we can get the class from the table name after we
     # retrieve something from the index
     # Dictionary of table name => table class.
     # Need the table name so we can get the class from the table name after we
     # retrieve something from the index


@@ -67,7 +106,7 @@ class PokedexLookup(object):
     )
 
 
     )
 
 

-    def __init__(self, directory=None, session=None, recreate=False):
+    def __init__(self, directory=None, session=None):

         """Opens the whoosh index stored in the named directory.  If the index
         doesn't already exist, it will be created.
 
         """Opens the whoosh index stored in the named directory.  If the index
         doesn't already exist, it will be created.
 


@@ -76,22 +115,19 @@ class PokedexLookup(object):
             `pokedex` egg directory.
 
         `session`
             `pokedex` egg directory.
 
         `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.
+            Used for creating the index and retrieving objects.  Defaults to an
+            attempt to connect to the default SQLite database installed by
+            `pokedex setup`.

         """
 
         # By the time this returns, self.index, self.speller, and self.session
         # must be set
 
         """
 
         # By the time this returns, self.index, self.speller, and self.session
         # must be set
 

-        # Defaults
-        if not directory:
-            directory = pkg_resources.resource_filename('pokedex',
-                                                        'data/whoosh-index')
+        # If a directory was not given, use the default
+        if directory is None:
+            directory = get_default_index_dir()
+
+        self.directory = directory

 
         if session:
             self.session = session
 
         if session:
             self.session = session


@@ -99,31 +135,34 @@ class PokedexLookup(object):
             self.session = connect()
 
         # Attempt to open or create the index
             self.session = connect()
 
         # 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!  Bam, done.
-            try:
-                self.index = whoosh.index.open_dir(directory, indexname='MAIN')
-                spell_store = whoosh.filedb.filestore.FileStorage(directory)
-                self.speller = whoosh.spelling.SpellChecker(spell_store)
-                return
-            except whoosh.index.EmptyIndexError as e:
-                # Apparently not a real index.  Fall out and create it
-                pass
-
-        # Delete and start over if we're going to bail anyway.
-        if directory_exists and recreate:
-            # Be safe and only delete if it looks like a whoosh index, i.e.,
-            # everything starts with _
-            if all(f[0] == '_' for f in os.listdir(directory)):
-                shutil.rmtree(directory)
-                directory_exists = False
-
-        if not directory_exists:
-            os.mkdir(directory)
-
-
-        ### Create index
+        if not os.path.exists(directory) or not os.listdir(directory):
+            # Directory doesn't exist OR is empty; caller needs to use
+            # rebuild_index before doing anything.  Provide a dummy object that
+            # complains when used
+            self.index = UninitializedIndex()
+            self.speller = UninitializedIndex()
+            return
+
+        # Otherwise, already exists; should be an index!  Bam, done.
+        # Note that this will explode if the directory exists but doesn't
+        # contain an index; that's a feature
+        try:
+            self.index = whoosh.index.open_dir(directory, indexname='MAIN')
+        except whoosh.index.EmptyIndexError:
+            raise IOError(
+                "The index directory already contains files.  "
+                "Please use a dedicated directory for the lookup index."
+            )
+
+        # Create speller, and done
+        spell_store = whoosh.filedb.filestore.FileStorage(directory)
+        self.speller = whoosh.spelling.SpellChecker(spell_store,
+            **self.SPELLER_OPTIONS)
+
+
+    def rebuild_index(self):
+        """Creates the index from scratch."""
+

         schema = whoosh.fields.Schema(
             name=whoosh.fields.ID(stored=True),
             table=whoosh.fields.ID(stored=True),
         schema = whoosh.fields.Schema(
             name=whoosh.fields.ID(stored=True),
             table=whoosh.fields.ID(stored=True),


@@ -133,16 +172,15 @@ class PokedexLookup(object):
             display_name=whoosh.fields.STORED,  # non-lowercased name
         )
 
             display_name=whoosh.fields.STORED,  # non-lowercased name
         )
 

-        self.index = whoosh.index.create_in(directory, schema=schema,
-                                            indexname='MAIN')
+        if not os.path.exists(self.directory):
+            os.mkdir(self.directory)
+
+        self.index = whoosh.index.create_in(self.directory, schema=schema,
+                                                            indexname='MAIN')

         writer = self.index.writer()
 
         # Index every name in all our tables of interest
         writer = self.index.writer()
 
         # Index every name in all our tables of interest

-        # speller_entries becomes a list of (word, score) tuples; the score is
-        # 2 for English names, 1.5 for Roomaji, and 1 for everything else.  I
-        # think this biases the results in the direction most people expect,
-        # especially when e.g. German names are very similar to English names
-        speller_entries = []
+        speller_entries = set()

         for cls in self.indexed_tables.values():
             q = self.session.query(cls)
 
         for cls in self.indexed_tables.values():
             q = self.session.query(cls)
 


@@ -150,7 +188,7 @@ class PokedexLookup(object):
                 row_key = dict(table=unicode(cls.__tablename__),
                                row_id=unicode(row.id))
 
                 row_key = dict(table=unicode(cls.__tablename__),
                                row_id=unicode(row.id))
 

-                def add(name, language, iso3166, score):
+                def add(name, language, iso3166):

                     normalized_name = self.normalize_name(name)
 
                     writer.add_document(
                     normalized_name = self.normalize_name(name)
 
                     writer.add_document(


@@ -159,21 +197,21 @@ class PokedexLookup(object):
                         **row_key
                     )
 
                         **row_key
                     )
 

-                    speller_entries.append((normalized_name, score))
+                    speller_entries.add(normalized_name)

 
 
                 # Add the basic English name to the index
                 if cls == tables.Pokemon:
                     # Pokémon need their form name added
                     # XXX kinda kludgy
 
 
                 # Add the basic English name to the index
                 if cls == tables.Pokemon:
                     # Pokémon need their form name added
                     # XXX kinda kludgy

-                    add(row.full_name, None, u'us', 1)
+                    add(row.full_name, None, u'us')

 
                     # If this is a default form, ALSO add the unadorned name,
                     # so 'Deoxys' alone will still do the right thing
                     if row.forme_name and not row.forme_base_pokemon_id:
 
                     # If this is a default form, ALSO add the unadorned name,
                     # so 'Deoxys' alone will still do the right thing
                     if row.forme_name and not row.forme_base_pokemon_id:

-                        add(row.name, None, u'us', 1)
+                        add(row.name, None, u'us')

                 else:
                 else:

-                    add(row.name, None, u'us', 1)
+                    add(row.name, None, u'us')

 
                 # Some things also have other languages' names
                 # XXX other language form names..?
 
                 # Some things also have other languages' names
                 # XXX other language form names..?


@@ -186,21 +224,21 @@ class PokedexLookup(object):
                         continue
 
                     add(moonspeak, foreign_name.language.name,
                         continue
 
                     add(moonspeak, foreign_name.language.name,

-                                   foreign_name.language.iso3166,
-                                   3)
+                                   foreign_name.language.iso3166)

 
                     # Add Roomaji too
                     if foreign_name.language.name == 'Japanese':
                         roomaji = romanize(foreign_name.name)
 
                     # Add Roomaji too
                     if foreign_name.language.name == 'Japanese':
                         roomaji = romanize(foreign_name.name)

-                        add(roomaji, u'Roomaji', u'jp', 8)
+                        add(roomaji, u'Roomaji', u'jp')

 
         writer.commit()
 
         # 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
 
         writer.commit()
 
         # 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

-        self.speller = whoosh.spelling.SpellChecker(self.index.storage)
-        self.speller.add_scored_words(speller_entries)
+        self.speller = whoosh.spelling.SpellChecker(self.index.storage, mingram=2,
+            **self.SPELLER_OPTIONS)
+        self.speller.add_words(speller_entries)

 
 
     def normalize_name(self, name):
 
 
     def normalize_name(self, name):


@@ -290,15 +328,17 @@ class PokedexLookup(object):
 
         Returns None for a bogus name.
         """
 
         Returns None for a bogus name.
         """

+        # Table object

         if hasattr(name, '__tablename__'):
             return getattr(name, '__tablename__')
         if hasattr(name, '__tablename__'):
             return getattr(name, '__tablename__')

-        elif name in self.indexed_tables:
-            return name
-        elif name + 's' in self.indexed_tables:
-            return name + 's'
-        else:
-            # Bogus.  Be nice and return dummy
-            return None
+
+        # Table name
+        for table in self.indexed_tables.values():
+            if name in (table.__tablename__, table.__singlename__):
+                return table.__tablename__
+
+        # Bogus.  Be nice and return dummy
+        return None

 
     def _whoosh_records_to_results(self, records, exact=True):
         """Converts a list of whoosh's indexed records to LookupResult tuples
 
     def _whoosh_records_to_results(self, records, exact=True):
         """Converts a list of whoosh's indexed records to LookupResult tuples


@@ -422,11 +462,30 @@ class PokedexLookup(object):
             exact = False
             results = []
 
             exact = False
             results = []
 

-            for suggestion in self.speller.suggest(
-                name, self.INTERMEDIATE_LOOKUP_RESULTS):
-
-                query = whoosh.query.Term('name', suggestion)
-                results.extend(searcher.search(query))
+            fuzzy_query_parts = []
+            fuzzy_weights = {}
+            min_weight = [None]
+            for suggestion, _, weight in self.speller.suggestions_and_scores(name):
+                # Only allow the top 50% of scores; otherwise there will always
+                # be a lot of trailing junk
+                if min_weight[0] is None:
+                    min_weight[0] = weight * 0.5
+                elif weight < min_weight[0]:
+                    break
+
+                fuzzy_query_parts.append(whoosh.query.Term('name', suggestion))
+                fuzzy_weights[suggestion] = weight
+
+            if not fuzzy_query_parts:
+                # Nothing at all; don't try querying
+                return []
+
+            fuzzy_query = whoosh.query.Or(fuzzy_query_parts)
+            if type_term:
+                fuzzy_query = fuzzy_query & type_term
+
+            searcher.weighting = LanguageWeighting(extra_weights=fuzzy_weights)
+            results = searcher.search(fuzzy_query)

 
         ### Convert results to db objects
         objects = self._whoosh_records_to_results(results, exact=exact)
 
         ### Convert results to db objects
         objects = self._whoosh_records_to_results(results, exact=exact)