Return a wrapper depending on whether an async or sync instance was instantiated

user_10211923 0
Viewed 603

Is there a way I can extract which subclass was instantiated given a method from a base class?

I know the question is a bit complicated, so here's an example:

from functools import wraps

def my_wrapper(fn_to_wrap):

    @wraps(fn_to_wrap)
    async def async_wrapper(*args, **kwargs):
        await do_some_async_stuff()
        print('Did some async stuff')
        return fn_to_wrap(*args, **kwargs)

    @wraps(fn_to_wrap)
    def sync_wrapper(*args, **kwargs):
        do some_sync_stuff()
        print('Did some sync stuff')
        return fn_to_wrap(*args, **kwargs)

    # <my_problem>
    if fn_to_wrap belongs_to(SyncClass):
        return sync_wrapper
    else:
        return async_wrapper
    # </my_problem>


class BaseClass:
    @my_wrapper
    def fn_to_wrap(self):
        return 'Finally a return a value'

class SyncClass(BaseClass):
    def fn_to_call(self):
        return self.fn_to_wrap()


class AsyncClass(BaseClass):
    async def fn_to_call(self):
        return await self.fn_to_wrap()

The problem is that the method fn_to_wrap belongs to the BaseClass. Which my Sync and Async classes inherit from.

Is there a way I can know whether fn_to_wrap belongs to an instance of AsyncClass or SyncClass?

Simply, I want my console to print:

>>> my_sync_class = SyncClass()
>>> print(my_sync_class.fn_to_call())

Done some sync stuff
Finally a return value

and 

>>> my_async_class = AsyncClass()
# not in a coroutine for brevity
>>> print(await my_async_class.fn_to_call())

Done some async stuff
FInally a return value

So, how would you implement </my_problem> to achieve those results?

[EDIT]

I'm aware of the existence of inspect.iscoroutinefunction and inspect.iscoroutine. But those won't help because the wrapped method is always synchronous while the wrapper is what does the async tasks.

async-await wrapper python-asyncio python-3.x python
1 Answers

If my_wrapper is allowed to know about AsyncClass and SyncClass (or you control them and can add a class attribute like _is_sync that tells the wrapper which kind of class it's dealing with), you can simply inspect self.

This cannot be done from the <my_problem> location because self is not yet available there; the code must return a single wrapper for both sync and async cases. Once called, the wrapper must detect the async case and return an instantiated async def if you need async behavior. (A sync function that returns a coroutine object is functionally equivalent to a coroutine function, much like an ordinary function that ends with return some_generator() is perfectly usable as a generator.)

Here is an example that uses isinstance to detect which variant is invoked:

def my_wrapper(fn_to_wrap):
    async def async_wrapper(*args, **kwargs):
        await asyncio.sleep(.1)
        print('Did some async stuff')
        return fn_to_wrap(*args, **kwargs)

    @wraps(fn_to_wrap)
    def uni_wrapper(self, *args, **kwargs):
        # or if self._is_async, etc.
        if isinstance(self, AsyncClass):
            return async_wrapper(self, *args, **kwargs)
        time.sleep(.1)
        print('Did some sync stuff')
        return fn_to_wrap(self, *args, **kwargs)

    return uni_wrapper

That implementation results in the desired output:

>>> x = SyncClass()
>>> x.fn_to_call()
Did some sync stuff
'Finally a return a value'
>>> async def test():
...     x = AsyncClass()
...     return await x.fn_to_call()
... 
>>> asyncio.get_event_loop().run_until_complete(test())
Did some async stuff
'Finally a return a value'

The above solution won't work if the wrapper cannot distinguish between SyncClass and AsyncClass. There are two constraints that might prevent it from doing so:

  • isinstance won't work if the number of subclasses is open-ended;
  • a custom class attribute won't work if the final classes are not controlled by the wrapper author.

In that case, the remaining option is to resort to black magic to determine whether the function is called from a coroutine or from a sync function. The black magic is conveniently provided by David Beazley in this talk:

def from_coroutine():
    return sys._getframe(2).f_code.co_flags & 0x380

Using from_coroutine, the uni_wrapper part of my_wrapper would look like this:

    @wraps(fn_to_wrap)
    def uni_wrapper(*args, **kwargs):
        if from_coroutine():
            return async_wrapper(*args, **kwargs)
        time.sleep(.1)
        print('Did some sync stuff')
        return fn_to_wrap(*args, **kwargs)

...providing the same result.

Of course, you have to be aware that the black magic might stop working in the next Python release without any warning whatsoever. But if you know what you're doing, it can come in quite useful.

user_16008980
Related