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.
This commit is contained in:
Eevee 2009-08-18 23:42:59 -07:00
parent fd5e863eed
commit 16072ceb44
2 changed files with 75 additions and 24 deletions

View file

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

View file

@ -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')
store = whoosh.filedb.filestore.RamStorage()
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
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 = 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