poros
10/4/2015 - 11:27 PM

Define a context manager usable as decorator

Define a context manager usable as decorator

# USAGE
@track_entry_and_exit('widget loader')
def activity():
    print('Some time consuming activity goes here')
    load_widget()

with track_entry_and_exit('widget loader'):
    print('Some time consuming activity goes here')
    load_widget()


# PYTHON 3
from contextlib import ContextDecorator

class track_entry_and_exit(ContextDecorator):
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        logging.info('Entering: {}'.format(self.name))

    def __exit__(self, exc_type, exc, exc_tb):
        logging.info('Exiting: {}'.format(self.name))


# PYTHON 2
class track_entry_and_exit(object):
    def __init__(self, name):
        self.name = name

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

    def __enter__(self):
        logging.info('Entering: {}'.format(self.name))

    def __exit__(self, err_type, value, traceback):
        logging.info('Exiting: {}'.format(self.name))
    

# MOCK.PATCH EXAMPLE
class testcase(object):
    def __init__(self):
        self.patcher = mock.patch('something', autospec=True)

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

    def __enter__(self):
        self.patcher.start()

    def __exit__(self, err_type, value, traceback):
        self.patcher.stop()

# note: @testcase(), even if without arguments