Search code examples
pythontwistedtornadocoroutine

Tornado context manager called in gen.coroutine


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.


Solution

  • 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))