Added setup command and made lookup work sanely. #15

The setup command loads the default data into a default location, then
creates a whoosh index in a default location.

get_index is now open_index and can be made to explicitly recreate the
index.  It also actually opens the index if it already existed, even
across processes, now that FileStorage is working.

The lookup command takes no switches for aiming at a different database;
it only uses the default data stores.

diff --git a/pokedex/__init__.py b/pokedex/__init__.py
index 6b788f7..1c28b06 100644 (file)
--- a/pokedex/__init__.py
+++ b/pokedex/__init__.py
@@ -4,7 +4,7 @@ import sys
 
 from .db import connect, metadata
 import pokedex.db.load
-from pokedex.lookup import lookup as pokedex_lookup
+import pokedex.lookup
 
 def main():
     if len(sys.argv) <= 1:
@@ -46,12 +46,14 @@ def command_load(*args):
                                   drop_tables=options.drop_tables,
                                   verbose=options.verbose)
 
+def command_setup(*args):
+    session = connect()
+    pokedex.db.load.load(session, verbose=False, drop_tables=True)
+    pokedex.lookup.open_index(session=session, recreate=True)
 
-def command_lookup(engine_uri, name):
-    # XXX don't require uri!  somehow
-    session = connect(engine_uri)
 
-    results, exact = pokedex_lookup(session, name)
+def command_lookup(name):
+    results, exact = pokedex.lookup.lookup(name)
     if exact:
         print "Matched:"
     else:
@@ -73,6 +75,9 @@ Commands:
 System commands:
     load                Load Pokédex data into a database from CSV files.
     dump                Dump Pokédex data from a database into CSV files.
+    setup               Loads Pokédex data into the right place and creates a
+                        lookup index in the right place.  No options or output.
+                        This will blow away the default database and index!
 
 Options:
     -d|--directory      By default, load and dump will use the CSV files in the

diff --git a/pokedex/lookup.py b/pokedex/lookup.py
index 2b38a1a..55e1b5a 100644 (file)
--- a/pokedex/lookup.py
+++ b/pokedex/lookup.py
@@ -1,5 +1,6 @@
 # encoding: utf8
 import os, os.path
+import pkg_resources
 import re
 
 from sqlalchemy.sql import func
@@ -10,6 +11,7 @@ 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.
@@ -36,17 +38,49 @@ extra_keys = {
     ],
 }
 
-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.
 
-    Creates an index if one does not exist.
+    `directory`
+        Directory containing the index.  Defaults to a location within the
+        `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.
     """
 
-    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()
+
+    # 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
 
-    store = whoosh.filedb.filestore.RamStorage()
+    if not directory_exists:
+        os.mkdir(directory)
+
+
+    # Create index
     schema = whoosh.fields.Schema(
         name=whoosh.fields.ID(stored=True),
         table=whoosh.fields.STORED,
@@ -58,10 +92,8 @@ def get_index(session):
         dummy=whoosh.fields.TEXT,
     )
 
-    index_directory = '/var/tmp/pokedex'
-    if not os.path.exists(index_directory):
-        os.mkdir(index_directory)
-    index = whoosh.index.create_in(index_directory, schema=schema)
+    index = whoosh.index.create_in(directory, schema=schema,
+                                              indexname='pokedex')
     writer = index.writer()
 
     # Index every name in all our tables of interest
@@ -103,7 +135,7 @@ def get_index(session):
     # 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,
@@ -125,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
@@ -142,11 +172,27 @@ 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 = open_index()
 
-    index, speller = get_index(session)
+    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