Context Manager/function/unittest class decorator hybrid

hybrid.py

import functools
import types


class fixit(object):
    """
    Context manager that can also be used as a function decorator or class decorator for unittest classes.
    """
    def __init__(self, **kwargs):
        self.kwargs = kwargs

    def __call__(self, func_or_class):
        if isinstance(func_or_class, (types.FunctionType, types.MethodType)):
            # function/method decorator
            func = func_or_class

            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                with self:
                    return func(*args, **kwargs)

            return wrapper
        elif isinstance(func_or_class, (types.TypeType, types.ClassType)) and hasattr(func_or_class, 'run'):
            # unittest class decorator
            klass = func_or_class
            klass.run = self(klass.run)
            return klass
        else:
            raise ValueError("Can only work with functions or unittest classes")

    def __enter__(self):
        print "Let's do it {name}".format(**self.kwargs)

    def __exit__(self, exc_type, exc_val, exc_tb):
        print "We've done it!"