Can a function be awaitable under the context of a running coroutine and be synchronous when called directly?

SargeATM 0
Viewed 46

I'm looking to create a function that is callable by both synchronous and asynchronous code. When called under the context of a running coroutine, it would return an awaitable object otherwise it would simply return the result.

Also I'm looking for a more general result that applies to using coroutines in all python not just asyncio.

import inspect

class awaitable:
    def __await__(self):
        return (yield 42)

def is_coroutine_running():
    #print(inspect.getouterframes(inspect.currentframe()))
    # ???
    return True 

def awaitable_when_called_by_coroutine():
    if is_coroutine_running():
        return awaitable()
    else:
        return 42

def sync_call():
    print("sync_call", awaitable_when_called_by_coroutine())

async def async_call():
    print("async_call", await awaitable_when_called_by_coroutine())

if __name__ == "__main__":
    sync_call()
    coro = async_call()
    assert(coro.send(None) == 42)
    try:
        coro.send(42)
    except StopIteration:
        pass

CURRENT OUTPUT

sync_call <__main__.awaitable object at 0x000002E0FE5347F0>
async_call 42

EXPECTED OUTPUT

sync_call 42
async_call 42

IDEAS

Search up the callstack to find a code object with inspect.CO_COROUTINE set

https://docs.python.org/3/library/inspect.html?highlight=inspect#inspect.CO_COROUTINE

See if CPython holds any state on the currently running coroutine.

coroutine python
2 Answers

Check for the inspect.CO_COROUTINE bit set on the code object of the outer frames

is_coroutine_running can be defined as follows:

def is_coroutine_running():
    frames = inspect.getouterframes(inspect.currentframe())
    return any(f.frame.f_code.co_flags & inspect.CO_COROUTINE for f in frames)
SargeATM0

Check for GET_AWAITABLE opcode

When awaiting a function call, a bytecode instruction GET_AWAITABLE is executed after CALL_FUNCTION. The next instruction in the caller's frame can be checked to see if it is GET_AWAITABLE.

EXAMPLE BYTECODE

 29           2 LOAD_GLOBAL              0 (print)
              4 LOAD_CONST               1 ('async_call')
              6 LOAD_GLOBAL              1 (awaitable_when_called_by_coroutine)
    -->       8 CALL_FUNCTION            0
             10 GET_AWAITABLE
             12 LOAD_CONST               0 (None)
             14 YIELD_FROM
             16 CALL_FUNCTION            2
             18 POP_TOP

SOLUTION

import inspect
import dis

class awaitable:
    def __await__(self):
        return (yield 42) 

def expects_awaitable():
    frames = inspect.getouterframes(inspect.currentframe())
    calling_frame = frames[2].frame
    bcode = dis.Bytecode(calling_frame.f_code)
    # Is the next instruction "GET_AWAITABLE"?
    for instr in bcode:
        if instr.offset > calling_frame.f_lasti:
            return instr.opname == 'GET_AWAITABLE'
    
    return False

def awaitable_when_called_by_coroutine():
    if expects_awaitable():
        return awaitable()
    else:
        return 42

def sync_call():
    print("sync_call", awaitable_when_called_by_coroutine())

async def async_call():
    print("async_call", await awaitable_when_called_by_coroutine())
    print("call sync_call in a coroutine")
    sync_call()

if __name__ == "__main__":
    sync_call()
    coro = async_call()
    assert(coro.send(None) == 42)
    try:
        coro.send(42)
    except StopIteration:
        pass
SargeATM0
Related