Skip to content

Commit

Permalink
Add async generator wrapper
Browse files Browse the repository at this point in the history
  • Loading branch information
TimPansino committed Jul 26, 2023
1 parent b9a91e5 commit fbe40ea
Show file tree
Hide file tree
Showing 4 changed files with 530 additions and 4 deletions.
51 changes: 47 additions & 4 deletions newrelic/common/async_wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,18 +18,23 @@
is_coroutine_callable,
is_asyncio_coroutine,
is_generator_function,
is_async_generator_function,
)

try:
import asyncio
except ImportError:
asyncio = None


def evaluate_wrapper(wrapper_string, wrapped, trace):
values = {'wrapper': None, 'wrapped': wrapped,
'trace': trace, 'functools': functools}
'trace': trace, 'functools': functools, "asyncio": asyncio}
exec(wrapper_string, values)
return values['wrapper']


def coroutine_wrapper(wrapped, trace):

WRAPPER = textwrap.dedent("""
@functools.wraps(wrapped)
async def wrapper(*args, **kwargs):
Expand All @@ -45,8 +50,6 @@ async def wrapper(*args, **kwargs):

def awaitable_generator_wrapper(wrapped, trace):
WRAPPER = textwrap.dedent("""
import asyncio
@functools.wraps(wrapped)
@asyncio.coroutine
def wrapper(*args, **kwargs):
Expand Down Expand Up @@ -81,9 +84,49 @@ def wrapper(*args, **kwargs):
return wrapper


def async_generator_wrapper(wrapped, trace):
WRAPPER = textwrap.dedent("""
@functools.wraps(wrapped)
async def wrapper(*args, **kwargs):
g = wrapped(*args, **kwargs)
value = None
with trace:
while True:
try:
g.asend(value).send(None)
except StopAsyncIteration as e:
# The underlying async generator has finished, return propagates a new StopAsyncIteration
return
except StopIteration as e:
# The call to async_generator_asend.send() should raise a StopIteration containing the yielded value
yielded = e.value
try:
value = yield yielded
except BaseException as e:
# An exception was thrown with .athrow(), propagate to the original async generator.
# Return value logic must be identical to .asend()
try:
g.athrow(type(e), e).send(None)
except StopAsyncIteration as e:
# The underlying async generator has finished, return propagates a new StopAsyncIteration
return
except StopIteration as e:
# The call to async_generator_athrow.send() should raise a StopIteration containing a yielded value
value = yield e.value
""")

try:
return evaluate_wrapper(WRAPPER, wrapped, trace)
except:
return wrapped


def async_wrapper(wrapped):
if is_coroutine_callable(wrapped):
return coroutine_wrapper
elif is_async_generator_function(wrapped):
return async_generator_wrapper
elif is_generator_function(wrapped):
if is_asyncio_coroutine(wrapped):
return awaitable_generator_wrapper
Expand Down
8 changes: 8 additions & 0 deletions newrelic/common/coroutine.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,3 +43,11 @@ def _iscoroutinefunction_tornado(fn):

def is_coroutine_callable(wrapped):
return is_coroutine_function(wrapped) or is_coroutine_function(getattr(wrapped, "__call__", None))


if hasattr(inspect, 'isasyncgenfunction'):
def is_async_generator_function(wrapped):
return inspect.isasyncgenfunction(wrapped)
else:
def is_async_generator_function(wrapped):
return False
Loading

0 comments on commit fbe40ea

Please sign in to comment.