contextdecorator.py

from __future__ import print_function

import abc
import functools


class BaseContextDecorator(object):
    __metaclass__ = abc.ABCMeta

    def __call__(self, func):

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

        return wrapped

    def __enter__(self):
        return self.start() or self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stop(exc_type, exc_val, exc_tb)

    @abc.abstractmethod
    def start(self):
        pass

    @abc.abstractmethod
    def stop(self, exc_type, exc_val, exc_tb):
        pass


if __name__ == '__main__':
    class Tag(BaseContextDecorator):
        def __init__(self, tag):
            self._tag = tag

        def start(self):
            print('<%s>' % self._tag)

        def stop(self, exc_type, exc_val, exc_tb):
            print('</%s>' % self._tag)

    def test_context_manager():
        print('Testing context manager')
        with Tag('h1'):
            print('Hello')

    def test_decorator():
        @Tag('h1')
        def h1(text):
            print(text)

        print('Testing decorator')
        print(h1('World'))

    test_context_manager()
    test_decorator()