Search code examples
python-3.xdecoratorcoroutinepython-3.4python-asyncio

How to decorate an asyncio.coroutine to retain its __name__?


I've tried to write a decorator function which wraps an asyncio.coroutine and returns the time it took to get done. The recipe below contains the code which is working as I expected. My only problem with it that somehow I loose the name of the decorated function despite the use of @functools.wraps. How to retain the name of the original coroutine? I checked the source of asyncio.

import asyncio
import functools
import random
import time

MULTIPLIER = 5

def time_resulted(coro):
    @functools.wraps(coro)
    @asyncio.coroutine
    def wrapper(*args, **kargs):
        time_before = time.time()
        result = yield from coro(*args, **kargs)
        if result is not None:
            raise TypeError('time resulted coroutine can '
                'only return None')
        return time_before, time.time()
    print('= wrapper.__name__: {!r} ='.format(wrapper.__name__))
    return wrapper

@time_resulted
@asyncio.coroutine
def random_sleep():
    sleep_time = random.random() * MULTIPLIER
    print('{} -> {}'.format(time.time(), sleep_time))
    yield from asyncio.sleep(sleep_time)

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    tasks = [asyncio.Task(random_sleep()) for i in range(5)]
    loop.run_until_complete(asyncio.wait(tasks))
    loop.close()
    for task in tasks:
        print(task, task.result()[1] - task.result()[0])
    print('= random_sleep.__name__: {!r} ='.format(
        random_sleep.__name__))
    print('= random_sleep().__name__: {!r} ='.format(
        random_sleep().__name__))

The result:

= wrapper.__name__: 'random_sleep' =
1397226479.00875 -> 4.261069174838891
1397226479.00875 -> 0.6596335046471768
1397226479.00875 -> 3.83421163259601
1397226479.00875 -> 2.5514027672929713
1397226479.00875 -> 4.497471439365472
Task(<wrapper>)<result=(1397226479.00875, 1397226483.274884)> 4.266134023666382
Task(<wrapper>)<result=(1397226479.00875, 1397226479.6697)> 0.6609499454498291
Task(<wrapper>)<result=(1397226479.00875, 1397226482.844265)> 3.835515022277832
Task(<wrapper>)<result=(1397226479.00875, 1397226481.562422)> 2.5536720752716064
Task(<wrapper>)<result=(1397226479.00875, 1397226483.51523)> 4.506479978561401
= random_sleep.__name__: 'random_sleep' =
= random_sleep().__name__: 'wrapper' =

As you can see random_sleep() returns a generator object with different name. I would like to retain the name of the decorated coroutine. I am not aware if this is problem is specific to asyncio.coroutines or not. I also tried the code with different decorator orders, but all has the same result. If I comment @functools.wraps(coro) then even random_sleep.__name__ becomes wrapper as I expected.

EDIT: I've posted this issue to Python Issue Tracker and received the following answer by R. David Murray: "I think this is a specific case of a more general need to improve 'wraps' that was discussed on python-dev not too long ago."


Solution

  • The issue is that functools.wraps changes only wrapper.__name__ and wrapper().__name__ stays wrapper. __name__ is a readonly generator attribute. You could use exec to set appropriate name:

    import asyncio
    import functools
    import uuid
    from textwrap import dedent
    
    def wrap_coroutine(coro, name_prefix='__' + uuid.uuid4().hex):
        """Like functools.wraps but preserves coroutine names."""
        # attribute __name__ is not writable for a generator, set it dynamically
        namespace = {
            # use name_prefix to avoid an accidental name conflict
            name_prefix + 'coro': coro,
            name_prefix + 'functools': functools,
            name_prefix + 'asyncio': asyncio,
        }
        exec(dedent('''
            def {0}decorator({0}wrapper_coro):
                @{0}functools.wraps({0}coro)
                @{0}asyncio.coroutine
                def {wrapper_name}(*{0}args, **{0}kwargs):
                    {0}result = yield from {0}wrapper_coro(*{0}args, **{0}kwargs)
                    return {0}result
                return {wrapper_name}
            ''').format(name_prefix, wrapper_name=coro.__name__), namespace)
        return namespace[name_prefix + 'decorator']
    

    Usage:

    def time_resulted(coro):
        @wrap_coroutine(coro)
        def wrapper(*args, **kargs):
            # ...
        return wrapper
    

    It works but there is probably a better way than using exec().