Require docutils.

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

diff --git a/pokedex/lookup.py b/pokedex/lookup.py
index b5b0c6d..28dd4b2 100644 (file)
--- a/pokedex/lookup.py
+++ b/pokedex/lookup.py
@@ -194,6 +194,31 @@ def _parse_table_name(name):
         # Bogus.  Be nice and return dummy
         return None
 
+def _whoosh_records_to_results(records, session, exact=True):
+    """Converts a list of whoosh's indexed records to LookupResult tuples
+    containing database objects.
+    """
+    # XXX this 'exact' thing is getting kinda leaky.  would like a better way
+    # to handle it, since only lookup() cares about fuzzy results
+    seen = {}
+    results = []
+    for record in records:
+        # Skip dupes
+        seen_key = record['table'], record['row_id']
+        if seen_key in seen:
+            continue
+        seen[seen_key] = True
+
+        cls = indexed_tables[record['table']]
+        obj = session.query(cls).get(record['row_id'])
+
+        results.append(LookupResult(object=obj,
+                                    name=record['display_name'],
+                                    language=record['language'],
+                                    exact=exact))
+
+    return results
+
 
 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.
@@ -206,7 +231,10 @@ def lookup(input, valid_types=[], session=None, indices=None, exact_only=False):
     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.
+
+    Extraneous whitespace is removed with extreme prejudice.
 
     Recognizes:
     - Names: "Eevee", "Surf", "Run Away", "Payapa Berry", etc.
@@ -251,18 +279,20 @@ def lookup(input, valid_types=[], session=None, indices=None, exact_only=False):
     else:
         index, speller = open_index()
 
-    name = unicode(input).lower()
+    name = unicode(input).strip().lower()
     exact = True
     form = None
 
     # Remove any type prefix (pokemon:133) before constructing a query
     if ':' in name:
-        prefix_chunk, name = name.split(':', 2)
-        prefixes = prefix_chunk.split(',')
+        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
-            valid_types = prefixes
+            prefixes = prefix_chunk.split(',')
+            valid_types = [_.strip() for _ in prefixes]
 
     # Random lookup
     if name == 'random':
@@ -287,7 +317,7 @@ def lookup(input, valid_types=[], session=None, indices=None, exact_only=False):
 
         # If there's a space in the input, this might be a form
         if ' ' in name:
-            form, formless_name = name.split(' ', 2)
+            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
@@ -319,22 +349,7 @@ def lookup(input, valid_types=[], session=None, indices=None, exact_only=False):
             results.extend(searcher.search(query))
 
     ### Convert results to db objects
-    objects = []
-    seen = {}
-    for result in results:
-        # Skip dupe results
-        seen_key = result['table'], result['row_id']
-        if seen_key in seen:
-            continue
-        seen[seen_key] = True
-
-        cls = indexed_tables[result['table']]
-        obj = session.query(cls).get(result['row_id'])
-
-        objects.append(LookupResult(object=obj,
-                                    name=result['display_name'],
-                                    language=result['language'],
-                                    exact=exact))
+    objects = _whoosh_records_to_results(results, session, 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
@@ -375,3 +390,28 @@ def random_lookup(valid_types=[], session=None, indices=None):
 
     return lookup(unicode(n), valid_types=[ partitions[0][0] ],
                   indices=indices, session=session)
+
+def prefix_lookup(prefix, session=None, indices=None):
+    """Returns terms starting with the given exact prefix.
+
+    No special magic is currently done with the name; type prefixes are not
+    recognized.
+
+    `session` and `indices` are treated as with `lookup()`.
+    """
+
+    if not session:
+        session = connect()
+
+    if indices:
+        index, speller = indices
+    else:
+        index, speller = open_index()
+
+    query = whoosh.query.Prefix(u'name', prefix.lower())
+
+    searcher = index.searcher()
+    searcher.weighting = LanguageWeighting()
+    results = searcher.search(query)  # XXX , limit=MAX_LOOKUP_RESULTS)
+
+    return _whoosh_records_to_results(results, session)