Tornado context manager called in gen.coroutine

Question

I am working now with AMQP pika library. What I want to develop is context manager or decorator that ensures connection establishment to Rabbit. The problem that it is impossible to use generators inside decorators because they must return called function. The following example raises an exception:

def ensure_conn(func):

    @gen.coroutine
    def wrapper(self, *args, **kwargs):
        yield self.connection.ensure_connection()
        return func(*args, **kwargs)

    return wrapper

There is almost the same problem with context managers. it is impossible to use yield two times. The following example raises exception that generator does not stopped.

@contextlib.contextmanager
@gen.coroutine
def ensure_conn(self):
    yield self.ensure_connection()
    yield

Please recommend another approach ? Of course I am already satisfied with simple coroutine's calls. Thank you.

Answer 1

Actually, there are two ways to create a context manager, that ensures something to you. In my case it was connection to AMQP. The first way is to override concurrent.futures.Future result() method, forcing it to return a generator function decorated by contextlib.contextmanager. ajdavis used this approach in his nice library TORO. You can see it by navigating to this line.

However, if you don't want to override concurrent.futures.Future object, then I recommend you walk through the following snippet:

@gen.coroutine
def ensure_connection(*args, **kwargs):
    res = yield _make_connection(timeout=kwargs.pop('timeout', 5), *args, **kwargs)
    raise gen.Return(res)

@gen.coroutine
def ensure(*args, **kwargs):
    res = yield ensure_connection(*args, **kwargs)

    @contextlib.contextmanager
    def func(res):
       try:
          yield  # your wrapped code
       finally:
          pass

    return func(res)

@gen.coroutine
def simple_usage(conn):
    with (yield conn.ensure()) as res:
        # Do your target staff here

conn = ...
IOLoop.add_callback(callback=lambda : simple_usage(conn))