helper for maintaining state across gevent greenlets

treelet.py

'''
A helper for gevent Greenlets -- tracks which greenlet created which,
and provides a handy place to stash data.

e.g.

g = treelet.spawn(foo)
g.locals['bar'] = 3
'''
import weakref
import sys

import gevent


class Treelet(gevent.Greenlet):
	def __init__(self, f, *a, **kw):
		super(Treelet, self).__init__(f, *a, **kw)
		spawner = self.spawn_parent = weakref.proxy(gevent.getcurrent())
		self.locals = getattr(spawner, 'locals', None) or {}
		stack = []
		cur = sys._getframe().f_back
		while cur:
			stack.extend((cur.f_code, cur.f_lineno))
			cur = cur.f_back
		self.stacks = (tuple(stack),) + getattr(spawner, 'stacks', ())[:10]


def spawn(f, *a, **kw):
	return Treelet.spawn(f, *a, **kw)



def test():
	import pprint

	def spawnN(n=4):
		if n == 0:
			pprint.pprint(gevent.getcurrent().stacks)
			return
		print gevent.getcurrent().locals.setdefault('count', 0)
		gevent.getcurrent().locals['count'] += 1
		gevent.wait([spawn(spawnN, n - 1)])
	gevent.wait([spawn(spawnN)])


	def make_stack(f, n=10):
		if n == 0:
			return f()
		return make_stack(f, n-1)

	def spawn_test():
		import time
		s = time.time()
		for i in range(int(1e3)):
			spawn(lambda: None)
		print "{0:0.2f} microseconds per spawn call".format(1000 * (time.time() - s))

	make_stack(spawn_test)