Monday, December 3, 2007

Python: asserting code runs in specific thread

My buddy JJ's post to the BayPiggies mailing list reminded me of a little snippet I wrote a while back that others might find useful as well. Personally, I avoid threads like the plague, but if you are forced to use them it is generally handy to keep accurate tabs on how you use them. In particular, as JJ suggested in his post, it is a good idea to assert that code is called in the thread context you expect it to be called in. This can go a long way toward avoiding one of many classes of hard-to-find logic bugs multi-threading invites. Anyway, on to the code...
def assertThread(*threads):
"""Decorator that asserts the wrapped function is only
run in a given thread
# If no thread list was supplied, assume the wrapped
# function should be run in the current thread.
if not threads:
threads = (threading.currentThread(), )

def decorator(func):
def wrapper(*args, **kw):
currentThread = threading.currentThread()
name = currentThread.getName()
assert currentThread in threads or name in threads, \
"%s erroniously called in %s thread " \
context" % (func.__name__, name)
return func(*args, **kw)

if __debug__:
return wrapper
return func
return decorator

You can restrict execution to one or more threads, each specified by either the thread object or thread name.

Note the trick at the end to make the decorator effectively a no-op in production. Using this decorator around your functions and methods helps you spot logic errors during development without impacting the performance of your production code. Of course, if you are of the school that assertions should never be disabled, feel free to replace the final if __debug__:/else block with an unconditional return of wrapper.