Refactor a bit to get better tracebacks (source lines).

Author: ionelmc

CHANGELOG.rst

@@ -2,6 +2,15 @@
 Changelog
 =========
 
+1.6.0 (2021-03-19)
+------------------
+
+* Added support for async special methods (``__aiter__``, ``__anext__``,
+  ``__await__``, ``__aenter__``, ``__aexit__``).
+  These are used in the ``async for``, ``await` and ``async with`` statements.
+
+  Note that ``__await__`` returns a wrapper that deals with the iterable/coroutine distinction
+
 
 1.5.2 (2020-11-26)
 ------------------

src/lazy_object_proxy/simple.py

@@ -258,20 +258,11 @@ def __reduce__(self):
     def __reduce_ex__(self, protocol):
         return identity, (self.__wrapped__,)
 
-    if await_ is not None:
-        exec("""
-def __aiter__(self):
-    return self.__wrapped__.__aiter__()
-
-async def __anext__(self):
-    return await self.__wrapped__.__anext__()
-
-def __await__(self):
-    return await_(self.__wrapped__)
-
-def __aenter__(self):
-    return self.__wrapped__.__aenter__()
-
-def __aexit__(self, *args, **kwargs):
-    return self.__wrapped__.__aexit__(*args, **kwargs)
-""")
+    if await_:
+        from .utils import __aenter__
+        from .utils import __aexit__
+        from .utils import __aiter__
+        from .utils import __anext__
+        from .utils import __await__
+
+        __aiter__, __anext__, __await__, __aenter__, __aexit__  # noqa

src/lazy_object_proxy/slots.py

@@ -429,20 +429,11 @@ def __reduce__(self):
     def __reduce_ex__(self, protocol):
         return identity, (self.__wrapped__,)
 
-    if await_ is not None:
-        exec("""
-def __aiter__(self):
-    return self.__wrapped__.__aiter__()
-
-async def __anext__(self):
-    return await self.__wrapped__.__anext__()
-
-def __await__(self):
-    return await_(self.__wrapped__)
-
-def __aenter__(self):
-    return self.__wrapped__.__aenter__()
-
-def __aexit__(self, *args, **kwargs):
-    return self.__wrapped__.__aexit__(*args, **kwargs)
-""")
+    if await_:
+        from .utils import __aenter__
+        from .utils import __aexit__
+        from .utils import __aiter__
+        from .utils import __anext__
+        from .utils import __await__
+
+        __aiter__, __anext__, __await__, __aenter__, __aexit__  # noqa

src/lazy_object_proxy/utils.py

@@ -1,3 +1,15 @@
+ # flake8: noqa
+try:
+    from .utils_py3 import __aenter__
+    from .utils_py3 import __aexit__
+    from .utils_py3 import __aiter__
+    from .utils_py3 import __anext__
+    from .utils_py3 import __await__
+    from .utils_py3 import await_
+except (ImportError, SyntaxError):
+    await_ = None
+
+
 def identity(obj):
     return obj
 
@@ -11,26 +23,3 @@ def __get__(self, obj, cls):
             return self
         value = obj.__dict__[self.func.__name__] = self.func(obj)
         return value
-
-
-try:
-    exec("""
-from inspect import isawaitable
-
-
-async def do_await(obj):
-    return await obj
-
-
-def do_yield_from(gen):
-    return (yield from gen)
-
-
-def await_(obj):
-    if isawaitable(obj):
-        return do_await(obj).__await__()
-    else:
-        return do_yield_from(obj)
-""")
-except (ImportError, SyntaxError):
-    await_ = None

src/lazy_object_proxy/utils_py3.py

@@ -0,0 +1,44 @@
+from collections.abc import Awaitable
+from inspect import CO_ITERABLE_COROUTINE
+from types import CoroutineType
+from types import GeneratorType
+
+
+async def do_await(obj):
+    return await obj
+
+
+def do_yield_from(gen):
+    return (yield from gen)
+
+
+def await_(obj):
+    obj_type = type(obj)
+    if (
+        obj_type is CoroutineType or
+        obj_type is GeneratorType and bool(obj.gi_code.co_flags & CO_ITERABLE_COROUTINE) or
+        isinstance(obj, Awaitable)
+    ):
+        return do_await(obj).__await__()
+    else:
+        return do_yield_from(obj)
+
+
+def __aiter__(self):
+    return self.__wrapped__.__aiter__()
+
+
+async def __anext__(self):
+    return await self.__wrapped__.__anext__()
+
+
+def __await__(self):
+    return await_(self.__wrapped__)
+
+
+def __aenter__(self):
+    return self.__wrapped__.__aenter__()
+
+
+def __aexit__(self, *args, **kwargs):
+    return self.__wrapped__.__aexit__(*args, **kwargs)

tests/test_async_py3.py

@@ -169,10 +169,7 @@ def bar():
         yield 2
 
     async def foo():
-        await proxy(lop.Proxy(bar))
-
-    import dis
-    dis.dis(foo)
+        await lop.Proxy(bar)
 
     f = lop.Proxy(foo)
     assert f.send(None) == 1