benchmark_docs.py

#!/usr/bin/env python

import cProfile
import pstats
import StringIO
import json

from spring.docgen import NewDocument, NewKey


def generate_keys():
    new_keys = NewKey(prefix=None, expiration=0)
    return tuple(new_keys.next(i)[0] for i in xrange(500000))

def run(keys):
    total_size = 0
    docs = NewDocument(avg_size=1024)
    for key in keys:
        d = docs.next(key)
        total_size += len(json.dumps(docs.next(key)))
    return total_size / len(keys)


def profile():
    pr = cProfile.Profile()
    s = StringIO.StringIO()

    keys = generate_keys()
    # print keys
    pr.enable()
    avg_size = run(keys)
    pr.disable()

    print "avg_size", avg_size

    ps = pstats.Stats(pr, stream=s).sort_stats('tottime')
    ps.reverse_order()
    ps.print_stats()
    ps.dump_stats('profile.prof')


if __name__ == '__main__':
    profile()