Fixed use of split() in lookup. Now strips whitespace. #94

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

diff --git a/pokedex/lookup.py b/pokedex/lookup.py
index 0bf1c18..1ba3fd1 100644 (file)
--- a/pokedex/lookup.py
+++ b/pokedex/lookup.py
@@ -2,7 +2,9 @@
 from collections import namedtuple
 import os, os.path
 import pkg_resources
 from collections import namedtuple
 import os, os.path
 import pkg_resources

+import random

 import re
 import re

+import shutil

 
 from sqlalchemy.sql import func
 import whoosh
 
 from sqlalchemy.sql import func
 import whoosh


@@ -10,12 +12,18 @@ import whoosh.filedb.filestore
 import whoosh.filedb.fileindex
 import whoosh.index
 from whoosh.qparser import QueryParser
 import whoosh.filedb.fileindex
 import whoosh.index
 from whoosh.qparser import QueryParser

+import whoosh.scoring

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

+__all__ = ['open_index', 'lookup', 'random_lookup']
+
+INTERMEDIATE_LOOKUP_RESULTS = 25
+MAX_LOOKUP_RESULTS = 10
+

 # 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


@@ -29,17 +37,6 @@ for cls in [
     ]:
     indexed_tables[cls.__tablename__] = cls
 
     ]:
     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),
-    ],
-}
-

 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.
 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.


@@ -61,7 +58,7 @@ def open_index(directory=None, session=None, recreate=False):
     # Defaults
     if not directory:
         directory = pkg_resources.resource_filename('pokedex',
     # Defaults
     if not directory:
         directory = pkg_resources.resource_filename('pokedex',

-                                                    'data/whoosh_index')
+                                                    'data/whoosh-index')

 
     if not session:
         session = connect()
 
     if not session:
         session = connect()


@@ -79,58 +76,74 @@ def open_index(directory=None, session=None, recreate=False):
             # Apparently not a real index.  Fall out of the if and create it
             pass
 
             # Apparently not a real index.  Fall out of the if 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)
 
 
     if not directory_exists:
         os.mkdir(directory)
 
 

-    # Create index
+    ### Create index

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

-        table=whoosh.fields.STORED,
-        row_id=whoosh.fields.STORED,
+        table=whoosh.fields.ID(stored=True),
+        row_id=whoosh.fields.ID(stored=True),

         language=whoosh.fields.STORED,
         language=whoosh.fields.STORED,

+        display_name=whoosh.fields.STORED,  # non-lowercased name
+        forme_name=whoosh.fields.ID,

     )
 
     index = whoosh.index.create_in(directory, schema=schema, indexname='MAIN')
     writer = index.writer()
 
     # Index every name in all our tables of interest
     )
 
     index = whoosh.index.create_in(directory, schema=schema, indexname='MAIN')
     writer = 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 = []
     for cls in indexed_tables.values():
         q = session.query(cls)
 
     speller_entries = []
     for cls in indexed_tables.values():
         q = session.query(cls)
 

-        # Only index base Pokémon formes
-        if hasattr(cls, 'forme_base_pokemon_id'):
-            q = q.filter_by(forme_base_pokemon_id=None)
-

         for row in q.yield_per(5):
         for row in q.yield_per(5):

-            row_key = dict(table=cls.__tablename__, row_id=row.id)
+            # XXX need to give forme_name a dummy value because I can't search
+            # for explicitly empty fields.  boo.
+            row_key = dict(table=unicode(cls.__tablename__),
+                           row_id=unicode(row.id),
+                           forme_name=u'XXX')
+
+            def add(name, language, score):
+                writer.add_document(name=name.lower(), display_name=name,
+                                    language=language,
+                                    **row_key)
+                speller_entries.append((name.lower(), score))

 
 

-            name = row.name.lower()
-            writer.add_document(name=name, **row_key)
-            speller_entries.append(name)
+            # If this is a form, mark it as such
+            if getattr(row, 'forme_base_pokemon_id', None):
+                row_key['forme_name'] = row.forme_name

 
 

-            for extra_key_func in extra_keys.get(cls, []):
-                extra_key = extra_key_func(row)
-                writer.add_document(name=extra_key, **row_key)
+            name = row.name
+            add(name, None, 1)

 
             # Pokemon also get other languages
 
             # Pokemon also get other languages

-            if cls == tables.Pokemon:
-                for foreign_name in row.foreign_names:
-                    name = foreign_name.name.lower()
-                    writer.add_document(name=name,
-                                        language=foreign_name.language.name,
-                                        **row_key)
-                    speller_entries.append(name)
-
-                    if foreign_name.language.name == 'Japanese':
-                        # Add Roomaji too
-                        roomaji = romanize(foreign_name.name).lower()
-                        writer.add_document(name=roomaji,
-                                            language='Roomaji',
-                                            **row_key)
-                        speller_entries.append(roomaji)
+            for foreign_name in getattr(row, 'foreign_names', []):
+                moonspeak = foreign_name.name
+                if name == moonspeak:
+                    # Don't add the English name again as a different language;
+                    # no point and it makes spell results confusing
+                    continue

 
 

+                add(moonspeak, foreign_name.language.name, 3)
+
+                # Add Roomaji too
+                if foreign_name.language.name == 'Japanese':
+                    roomaji = romanize(foreign_name.name)
+                    add(roomaji, u'Roomaji', 8)

 
     writer.commit()
 
 
     writer.commit()
 


@@ -138,30 +151,86 @@ def open_index(directory=None, session=None, recreate=False):
     # at once, as every call to add_* does a commit(), and those seem to be
     # expensive
     speller = whoosh.spelling.SpellChecker(index.storage)
     # at once, as every call to add_* does a commit(), and those seem to be
     # expensive
     speller = whoosh.spelling.SpellChecker(index.storage)

-    speller.add_words(speller_entries)
+    speller.add_scored_words(speller_entries)

 
     return index, speller
 
 
 
     return index, speller
 
 

-LookupResult = namedtuple('LookupResult', ['object', 'language', 'exact'])
-def lookup(name, session=None, indices=None, exact_only=False):
+class LanguageWeighting(whoosh.scoring.Weighting):
+    """A scoring class that forces otherwise-equal English results to come
+    before foreign results.
+    """
+
+    def score(self, searcher, fieldnum, text, docnum, weight, QTF=1):
+        doc = searcher.stored_fields(docnum)
+        if doc['language'] == None:
+            # English (well, "default"); leave it at 1
+            return weight
+        elif doc['language'] == u'Roomaji':
+            # Give Roomaji a bit of a boost, as it's most likely to be searched
+            return weight * 0.95
+        else:
+            # Everything else can drop down the totem pole
+            return weight * 0.9
+
+rx_is_number = re.compile('^\d+$')
+
+LookupResult = namedtuple('LookupResult',
+                          ['object', 'name', 'language', 'exact'])
+
+def _parse_table_name(name):
+    """Takes a singular table name, table name, or table object and returns the
+    table name.
+
+    Returns None for a bogus name.
+    """
+    if hasattr(name, '__tablename__'):
+        return getattr(name, '__tablename__')
+    elif name in indexed_tables:
+        return name
+    elif name + 's' in indexed_tables:
+        return name + 's'
+    else:
+        # Bogus.  Be nice and return dummy
+        return None
+
+
+def lookup(input, valid_types=[], session=None, indices=None, exact_only=False):

     """Attempts to find some sort of object, given a database session and name.
 
     """Attempts to find some sort of object, given a database session and name.
 

-    Returns a list of named (object, language, exact) tuples.  `object` is a
-    database object, `language` is the name of the language in which the name
-    was found, and `exact` is True iff this was an exact match.
+    Returns a list of named (object, name, language, exact) tuples.  `object`
+    is a database object, `name` is the name under which the object was found,
+    `language` is the name of the language in which the name was found, and
+    `exact` is True iff this was an exact match.

 
     This function currently ONLY does fuzzy matching if there are no exact
     matches.
 
 
     This function currently ONLY does fuzzy matching if there are no exact
     matches.
 

-    Formes are not returned; "Shaymin" will return only grass Shaymin.
+    Formes are not returned unless requested; "Shaymin" will return only grass
+    Shaymin.

 
 

-    Recognizes:
-    - Pokémon names: "Eevee"
+    Extraneous whitespace is removed with extreme prejudice.

 
 

-    `name`
+    Recognizes:
+    - Names: "Eevee", "Surf", "Run Away", "Payapa Berry", etc.
+    - Foreign names: "Iibui", "Eivui"
+    - Fuzzy names in whatever language: "Evee", "Ibui"
+    - IDs: "133", "192", "250"
+    Also:
+    - Type restrictions.  "type:psychic" will only return the type.  This is
+      how to make ID lookup useful.  Multiple type specs can be entered with
+      commas, as "move,item:1".  If `valid_types` are provided, any type prefix
+      will be ignored.
+    - Alternate formes can be specified merely like "wash rotom".
+
+    `input`

         Name of the thing to look for.
 
         Name of the thing to look for.
 

+    `valid_types`
+        A list of table objects or names, e.g., `['pokemon', 'moves']`.  If
+        this is provided, only results in one of the given tables will be
+        returned.
+

     `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
     `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


@@ -185,22 +254,72 @@ def lookup(name, session=None, indices=None, exact_only=False):
     else:
         index, speller = open_index()
 
     else:
         index, speller = open_index()
 

-    name = unicode(name)
-
+    name = unicode(input).strip().lower()

     exact = True
     exact = True

+    form = None
+
+    # Remove any type prefix (pokemon:133) before constructing a query
+    if ':' in name:
+        prefix_chunk, name = name.split(':', 1)
+        name = name.strip()
+
+        if not valid_types:
+            # Only use types from the query string if none were explicitly
+            # provided
+            prefixes = prefix_chunk.split(',')
+            valid_types = [_.strip() for _ in prefixes]
+
+    # Random lookup
+    if name == 'random':
+        return random_lookup(indices=(index, speller),
+                             session=session,
+                             valid_types=valid_types)
+
+    # Do different things depending what the query looks like
+    # Note: Term objects do an exact match, so we don't have to worry about a
+    # query parser tripping on weird characters in the input
+    if '*' in name or '?' in name:
+        exact_only = True
+        query = whoosh.query.Wildcard(u'name', name)
+    elif rx_is_number.match(name):
+        # Don't spell-check numbers!
+        exact_only = True
+        query = whoosh.query.Term(u'row_id', name)
+    else:
+        # Not an integer
+        query = whoosh.query.Term(u'name', name) \
+              & whoosh.query.Term(u'forme_name', u'XXX')
+
+        # If there's a space in the input, this might be a form
+        if ' ' in name:
+            form, formless_name = name.split(' ', 1)
+            form_query = whoosh.query.Term(u'name', formless_name) \
+                       & whoosh.query.Term(u'forme_name', form)
+            query = query | form_query
+
+    ### Filter by type of object
+    type_terms = []
+    for valid_type in valid_types:
+        table_name = _parse_table_name(valid_type)
+        type_terms.append(whoosh.query.Term(u'table', table_name))

 
 

-    # 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
+    if type_terms:
+        query = query & whoosh.query.Or(type_terms)
+
+
+    ### Actual searching

     searcher = index.searcher()
     searcher = index.searcher()

-    query = whoosh.query.Term('name', name.lower())
-    results = searcher.search(query)
+    searcher.weighting = LanguageWeighting()  # XXX kosher?  docs say search()
+                                              # takes a weighting kw but it
+                                              # certainly does not
+    results = searcher.search(query, limit=INTERMEDIATE_LOOKUP_RESULTS)

 
     # Look for some fuzzy matches if necessary
     if not exact_only and not results:
         exact = False
         results = []
 
 
     # Look for some fuzzy matches if necessary
     if not exact_only and not results:
         exact = False
         results = []
 

-        for suggestion in speller.suggest(name, 10):
+        for suggestion in speller.suggest(name, INTERMEDIATE_LOOKUP_RESULTS):

             query = whoosh.query.Term('name', suggestion)
             results.extend(searcher.search(query))
 
             query = whoosh.query.Term('name', suggestion)
             results.extend(searcher.search(query))
 


@@ -216,6 +335,48 @@ def lookup(name, session=None, indices=None, exact_only=False):
 
         cls = indexed_tables[result['table']]
         obj = session.query(cls).get(result['row_id'])
 
         cls = indexed_tables[result['table']]
         obj = session.query(cls).get(result['row_id'])

-        objects.append(LookupResult(obj, result['language'], exact))

 
 

-    return objects
+        objects.append(LookupResult(object=obj,
+                                    name=result['display_name'],
+                                    language=result['language'],
+                                    exact=exact))
+
+    # Only return up to 10 matches; beyond that, something is wrong.
+    # We strip out duplicate entries above, so it's remotely possible that we
+    # should have more than 10 here and lost a few.  The speller returns 25 to
+    # give us some padding, and should avoid that problem.  Not a big deal if
+    # we lose the 25th-most-likely match anyway.
+    return objects[:MAX_LOOKUP_RESULTS]
+
+
+def random_lookup(valid_types=[], session=None, indices=None):
+    """Takes similar arguments as `lookup()`, but returns a random lookup
+    result from one of the provided `valid_types`.
+    """
+
+    tables = []
+    for valid_type in valid_types:
+        table_name = _parse_table_name(valid_type)
+        if table_name:
+            tables.append(indexed_tables[table_name])
+
+    if not tables:
+        tables = indexed_tables.values()
+
+    # Rather than create an array of many hundred items and pick randomly from
+    # it, just pick a number up to the total number of potential items, then
+    # pick randomly from that, and partition the whole range into chunks
+    total = 0
+    partitions = []
+    for table in tables:
+        count = session.query(table).count()
+        total += count
+        partitions.append((table, count))
+
+    n = random.randint(1, total)
+    while n > partitions[0][1]:
+        n -= partitions[0][1]
+        partitions.pop(0)
+
+    return lookup(unicode(n), valid_types=[ partitions[0][0] ],
+                  indices=indices, session=session)