[Python-ideas] Fast context creation

I have a situation where, no matter how the routine ends, I need to return the focus to a widget. It would be rather clever to do this with a context, but I would expect something like this to be a possible strategy with contextlib.context(enter=None, exit=lambda *args: my_widget.setFocus()): do what I need to do as far as I know, at the moment it's not possible. Am I right? I think it would be an easy and practical addition to the contextlib module to quickly register two routines for enter and exit.