autocomplete.py

"""
A redis autocomplete example for multi-word phrases.

Based on:
  Ruby original: http://gist.github.com/574044
  Python original: https://gist.github.com/577852

See options below for usage

Requires http://github.com/andymccurdy/redis-py/
"""

from redis import Redis
import sys

r = Redis()
ZKEY_COMPL = 'compl'
SKEY_DOCS_PREFIX = 'docs:'

def deleteAll():
    """Clear out the completions db"""
    return r.zremrangebyrank(ZKEY_COMPL, 0, -1)

def addCompletions(text):
    """Create the completion sorted set."""
    text = text.strip()
    if not text:
        return

    for word in text.split():
        word = word.lower()
        for end_index in range(1, len(word)+1):
            prefix = word[:end_index]
            r.zadd(ZKEY_COMPL, prefix, 0)
        r.zadd(ZKEY_COMPL, word + '*', 0)     
        r.sadd(SKEY_DOCS_PREFIX + word, text)

    r.zadd(ZKEY_COMPL, text + '*', 0)
    r.sadd(SKEY_DOCS_PREFIX + text.lower(), text)

def addFromFile(filename):
    """Create completions for all lines in filename"""
    print "Adding completions for", filename, "...",
    sys.stdout.flush()
    for line in open(filename):
        addCompletions(line)
    print "done"

def getWordCompletions(r, word, count, rangelen=50):
    """Get up to count completions for the given word"""
    prefix = word.lower().strip()
    results = set()
    if not prefix:
        return results

    start = r.zrank(ZKEY_COMPL, prefix)
    if not start:
        return results

    while len(results) <= count:
        entries = r.zrange(ZKEY_COMPL, start, start + rangelen - 1)
        start += rangelen

        if not entries or len(entries) == 0:
            break

        for entry in entries:
            minlen = min((len(entry), len(prefix)))
            if entry[:minlen] != prefix[:minlen]:
                return results

            if entry[-1] == '*' and len(results) <= count:
                results.add(entry[0:-1])

    return results

def getPhraseCompletions(r, text, count):
    """
    Get up to @count completions for @text

    For an input text of N words, uses N+1 redis calls: 
    One ZRANK per word, and one SUNION at the end.
    """
    results = set()
    for prefix in text.lower().split():
        results.update(getWordCompletions(r, prefix, count))

    keys = map(lambda k: SKEY_DOCS_PREFIX+k, results) 
    if keys:
        return sorted(r.sunion( keys ), key=str.lower)[:count]
    else:
        return []

if __name__ == '__main__':
    from optparse import OptionParser
    parser = OptionParser()

    parser.add_option("-f", "--file", dest="filename",
                      help="Create completions for lines in FILE", metavar="FILE")
    parser.add_option("-i", "--insert", dest="text",
                      help="Create completions for TEXT", metavar="TEXT")
    parser.add_option("-d", "--delete-all", dest="delete", 
                      action="store_true", default=False,
                      help="Delete everything in the completions db")

    (options, args) = parser.parse_args()
    
    if options.delete:
        deleteAll()

    if options.filename:
        addFromFile(options.filename)

    if options.text:
        addCompletions(options.text)

    if args:
        for arg in args:
            print arg, '==>'
            print '\n'.join(getPhraseCompletions(r, arg, 10))