Merge pull request #78 from broglep-work/support-asyncgen

Add support for async generator functions

Author: smarie

src/makefun/_main_latest_py.py

@@ -22,3 +22,22 @@ def partial_f(*args, **kwargs):
         kwargs.update(preset_kwargs)  # for python 3.4: explicit dict update
         yield from f(*chain(preset_pos_args, args), **kwargs)
     return partial_f
+
+
+def make_partial_using_async_for_in_yield(new_sig, f, *preset_pos_args, **preset_kwargs):
+    """
+    Makes a 'partial' when f is a async generator and python is new enough to support `async for v in f(): yield v`
+
+    :param new_sig:
+    :param f:
+    :param presets:
+    :return:
+    """
+
+    @wraps(f, new_sig=new_sig)
+    async def partial_f(*args, **kwargs):
+        kwargs.update(preset_kwargs)
+        async for v in f(*chain(preset_pos_args, args), **kwargs):
+            yield v
+
+    return partial_f

src/makefun/main.py

@@ -33,6 +33,13 @@ def iscoroutinefunction(f):
     def isgeneratorfunction(f):
         return False
 
+try:
+    from inspect import isasyncgenfunction
+except ImportError:
+    # assume no generator function in old Python versions
+    def isasyncgenfunction(f):
+        return False
+
 try:  # python 3.5+
     from typing import Callable, Any, Union, Iterable, Dict, Tuple, Mapping
 except ImportError:
@@ -246,6 +253,8 @@ def create_function(func_signature,             # type: Union[str, Signature]
         else:
             from makefun._main_legacy_py import get_legacy_py_generator_body_template
             body = get_legacy_py_generator_body_template() % (func_signature_str, params_str)
+    elif isasyncgenfunction(func_impl):
+        body = "async def %s\n    async for y in _func_impl_(%s):\n        yield y\n" % (func_signature_str, params_str)
     else:
         body = "def %s\n    return _func_impl_(%s)\n" % (func_signature_str, params_str)
 
@@ -1123,6 +1132,9 @@ def partial(f,                 # type: Callable
         else:
             from makefun._main_legacy_py import make_partial_using_yield
             partial_f = make_partial_using_yield(new_sig, f, *preset_pos_args, **preset_kwargs)
+    elif isasyncgenfunction(f) and sys.version_info >= (3, 6):
+        from makefun._main_latest_py import make_partial_using_async_for_in_yield
+        partial_f = make_partial_using_async_for_in_yield(new_sig, f, *preset_pos_args, **preset_kwargs)
     else:
         @wraps(f, new_sig=new_sig)
         def partial_f(*args, **kwargs):

tests/_test_py36.py

@@ -0,0 +1,21 @@
+from makefun import wraps, partial
+
+
+def make_async_generator():
+    async def f(v):
+        yield v
+
+    return f
+
+
+def make_async_generator_wrapper(f):
+    @wraps(f)
+    async def wrapper(*args, **kwargs):
+        async for v in f(*args, **kwargs):
+            yield v
+
+    return wrapper
+
+
+def make_async_generator_partial(f, *args, **kwargs):
+    return partial(f, *args, **kwargs)

tests/test_issues.py

@@ -1,3 +1,4 @@
+import inspect
 import sys
 
 import pytest
@@ -224,3 +225,31 @@ def f(a):
 
     f2 = create_function("zoo(a)", f, func=f)
     assert f2(3) == 4
+
+
+@pytest.mark.skipif(sys.version_info < (3, 6), reason="requires python 3.6 or higher (async generator)")
+def test_issue_77_async_generator_wraps():
+    import asyncio
+    from ._test_py36 import make_async_generator, make_async_generator_wrapper
+
+    f = make_async_generator()
+    wrapper = make_async_generator_wrapper(f)
+
+    assert inspect.isasyncgenfunction(f)
+    assert inspect.isasyncgenfunction(wrapper)
+
+    assert asyncio.get_event_loop().run_until_complete(asyncio.ensure_future(wrapper(1).__anext__())) == 1
+
+
+@pytest.mark.skipif(sys.version_info < (3, 6), reason="requires python 3.6 or higher (async generator)")
+def test_issue_77_async_generator_partial():
+    import asyncio
+    from ._test_py36 import make_async_generator, make_async_generator_partial
+
+    f = make_async_generator()
+    f_partial = make_async_generator_partial(f, v=1)
+
+    assert inspect.isasyncgenfunction(f)
+    assert inspect.isasyncgenfunction(f_partial)
+
+    assert asyncio.get_event_loop().run_until_complete(asyncio.ensure_future(f_partial().__anext__())) == 1