trie.py

import datrie, random, geohash, guppy, time, itertools

ghashBase32Chars = '0123456789bcdefghjkmnpqrstuvwxyz'

def timeme(func):
    def wrapper(*arg,**kw):
        t1 = time.time()
        res = func(*arg,**kw)
        t2 = time.time()
        return (t2-t1), res
    return wrapper

def randomHash():
    lat = (random.random() * 180) - 90
    lng = (random.random() * 360) - 180
    return geohash.encode(lat, lng)

def insert(t, hash, userId):
    t.setdefault(hash, set()).add(userId)

if __name__ == '__main__':
    h = guppy.hpy()
    h.setrelheap()

    trie = datrie.Trie(ghashBase32Chars)
    ordinals = itertools.count()
    for ordinal in xrange(50000):
        hash = randomHash().decode('ascii')
        insert(trie, hash, ordinals.next())

    print 'Memory usage for 50k hashes'
    print h.heap()

    @timeme
    def insertMany(hashes):
        for hash in hashes:
            insert(trie, hash, ordinals.next())
    hashes = [randomHash().decode('ascii') for _ in xrange(100000)]
    (msecs, _) = insertMany(hashes)
    print '100,000 inserts: %.2fs' % (msecs,)

    @timeme
    def lookupMany(hashes):
        for hash in hashes:
            trie[hash]
    (msecs, _) = lookupMany(hashes)
    print '100,000 lookups: %.2fs' % (msecs,)