Python: My simple ContextManager implementation

context_manager.py

import time


class ContextManager():
    def __init__(self, generator_method, args, kwargs):
        self.generator_method = generator_method
        self.args = args
        self.kwargs = kwargs

    def __enter__(self):
        self.generator = self.__run_generator()
        next(self.generator)

    def __exit__(self, *_):
        for _ in self.generator:
            pass

    def __run_generator(self):
        return self.generator_method(*self.args, **self.kwargs)


def context_manager(generator_method):
    def method_with_context(*args, **kwargs):
        return ContextManager(generator_method, args, kwargs)

    return method_with_context


@context_manager
def measure(task_name):
    t = time.time()
    yield
    print(task_name, 'took', time.time() - t, 'seconds.')


with measure('sleep test') as method:
    time.sleep(5)