Add support for async generators to "@logger.catch" (#1303)

The implementation is ugly because of Python 3.5 required compatibility
("yield" inside "async def" would generate "SyntaxError").

Also, note that currently async generator do not support arbitrary
return values, but this might change (there is an opened PR at the time
of writing).

Author: Delgan

CHANGELOG.rst

@@ -7,6 +7,8 @@
 - Update the default log format to include the timezone offset since it produces less ambiguous logs (`#856 <https://github.com/Delgan/loguru/pull/856>`_, thanks `@tim-x-y-z <https://github.com/tim-x-y-z>`_).
 - Honor the ``NO_COLOR`` environment variable to disable color output by default if ``colorize`` is not provided (`#1178 <https://github.com/Delgan/loguru/issues/1178>`_).
 - Add requirement for ``TERM`` environment variable not to be ``"dumb"`` to enable colorization (`#1287 <https://github.com/Delgan/loguru/pull/1287>`_, thanks `@snosov1 <https://github.com/snosov1>`_).
+- Make ``logger.catch()`` compatible with asynchronous generators (`#1302 <https://github.com/Delgan/loguru/issues/1302>`_).
+
 
 `0.7.3`_ (2024-12-06)
 =====================

loguru/_logger.py

@@ -114,10 +114,16 @@
 from ._simple_sinks import AsyncSink, CallableSink, StandardSink, StreamSink
 
 if sys.version_info >= (3, 6):
+    from collections.abc import AsyncGenerator
+    from inspect import isasyncgenfunction
     from os import PathLike
+
 else:
     from pathlib import PurePath as PathLike
 
+    def isasyncgenfunction(func):
+        return False
+
 
 Level = namedtuple("Level", ["name", "no", "color", "icon"])  # noqa: PYI024
 
@@ -1293,6 +1299,30 @@ def catch_wrapper(*args, **kwargs):
                             return (yield from function(*args, **kwargs))
                         return default
 
+                elif isasyncgenfunction(function):
+
+                    class AsyncGenCatchWrapper(AsyncGenerator):
+
+                        def __init__(self, gen):
+                            self._gen = gen
+
+                        async def asend(self, value):
+                            with catcher:
+                                try:
+                                    return await self._gen.asend(value)
+                                except StopAsyncIteration:
+                                    pass
+                                except:
+                                    raise
+                            raise StopAsyncIteration
+
+                        async def athrow(self, *args, **kwargs):
+                            return await self._gen.athrow(*args, **kwargs)
+
+                    def catch_wrapper(*args, **kwargs):
+                        gen = function(*args, **kwargs)
+                        return AsyncGenCatchWrapper(gen)
+
                 else:
 
                     def catch_wrapper(*args, **kwargs):

pyproject.toml

@@ -133,6 +133,9 @@ convention = "numpy"
 [tool.typos.default]
 extend-ignore-re = ["(?Rm)^.*# spellchecker: disable-line$"]
 
+[tool.typos.default.extend-identifiers]
+asend = "asend"
+
 [tool.typos.files]
 extend-exclude = [
   "tests/exceptions/output/**", # False positive due to ansi sequences.

tests/exceptions/output/modern/decorate_async_generator.txt

@@ -0,0 +1 @@
+Done

tests/exceptions/output/modern/exception_formatting_async_generator.txt

@@ -0,0 +1,42 @@
+
+Traceback (most recent call last):
+  File "tests/exceptions/source/modern/exception_formatting_async_generator.py", line 20, in <module>
+    f.send(None)
+  File "tests/exceptions/source/modern/exception_formatting_async_generator.py", line 14, in foo
+    yield a / b
+ZeroDivisionError: division by zero
+
+Traceback (most recent call last):
+
+  File "tests/exceptions/source/modern/exception_formatting_async_generator.py", line 20, in <module>
+    f.send(None)
+    │ └ <method 'send' of 'coroutine' objects>
+    └ <coroutine object Logger.catch.<locals>.Catcher.__call__.<locals>.AsyncGenCatchWrapper.asend at 0xDEADBEEF>
+
+  File "tests/exceptions/source/modern/exception_formatting_async_generator.py", line 14, in foo
+    yield a / b
+          │   └ 0
+          └ 1
+
+ZeroDivisionError: division by zero
+
+Traceback (most recent call last):
+> File "tests/exceptions/source/modern/exception_formatting_async_generator.py", line 20, in <module>
+    f.send(None)
+  File "tests/exceptions/source/modern/exception_formatting_async_generator.py", line 14, in foo
+    yield a / b
+ZeroDivisionError: division by zero
+
+Traceback (most recent call last):
+
+> File "tests/exceptions/source/modern/exception_formatting_async_generator.py", line 20, in <module>
+    f.send(None)
+    │ └ <method 'send' of 'coroutine' objects>
+    └ <coroutine object Logger.catch.<locals>.Catcher.__call__.<locals>.AsyncGenCatchWrapper.asend at 0xDEADBEEF>
+
+  File "tests/exceptions/source/modern/exception_formatting_async_generator.py", line 14, in foo
+    yield a / b
+          │   └ 0
+          └ 1
+
+ZeroDivisionError: division by zero

tests/exceptions/source/modern/decorate_async_generator.py

@@ -0,0 +1,111 @@
+from loguru import logger
+import asyncio
+import sys
+
+logger.remove()
+
+# We're truly only testing whether the tests succeed, we do not care about the formatting.
+# These should be regular Pytest test cases, but that is not possible because the syntax is not valid in Python 3.5.
+logger.add(lambda m: None, format="", diagnose=True, backtrace=True, colorize=True)
+
+def test_decorate_async_generator():
+    @logger.catch(reraise=True)
+    async def generator(x, y):
+        yield x
+        yield y
+
+    async def coro():
+        out = []
+        async for val in generator(1, 2):
+            out.append(val)
+        return out
+
+    res = asyncio.run(coro())
+    assert res == [1, 2]
+
+
+def test_decorate_async_generator_with_error():
+    @logger.catch(reraise=False)
+    async def generator(x, y):
+        yield x
+        yield y
+        raise ValueError
+
+    async def coro():
+        out = []
+        async for val in generator(1, 2):
+            out.append(val)
+        return out
+
+    res = asyncio.run(coro())
+    assert res == [1, 2]
+
+def test_decorate_async_generator_with_error_reraised():
+    @logger.catch(reraise=True)
+    async def generator(x, y):
+        yield x
+        yield y
+        raise ValueError
+
+    async def coro():
+        out = []
+        try:
+            async for val in generator(1, 2):
+                out.append(val)
+        except ValueError:
+            pass
+        else:
+            raise AssertionError("ValueError not raised")
+        return out
+
+    res = asyncio.run(coro())
+    assert res == [1, 2]
+
+
+def test_decorate_async_generator_then_async_send():
+    @logger.catch
+    async def generator(x, y):
+        yield x
+        yield y
+
+    async def coro():
+        gen = generator(1, 2)
+        await gen.asend(None)
+        await gen.asend(None)
+        try:
+            await gen.asend(None)
+        except StopAsyncIteration:
+            pass
+        else:
+            raise AssertionError("StopAsyncIteration not raised")
+
+    asyncio.run(coro())
+
+
+def test_decorate_async_generator_then_async_throw():
+    @logger.catch
+    async def generator(x, y):
+        yield x
+        yield y
+
+    async def coro():
+        gen = generator(1, 2)
+        await gen.asend(None)
+        try:
+            await gen.athrow(ValueError)
+        except ValueError:
+            pass
+        else:
+            raise AssertionError("ValueError not raised")
+
+    asyncio.run(coro())
+
+
+test_decorate_async_generator()
+test_decorate_async_generator_with_error()
+test_decorate_async_generator_with_error_reraised()
+test_decorate_async_generator_then_async_send()
+test_decorate_async_generator_then_async_throw()
+
+logger.add(sys.stderr, format="{message}")
+logger.info("Done")

tests/exceptions/source/modern/exception_formatting_async_generator.py

@@ -0,0 +1,22 @@
+import sys
+
+from loguru import logger
+
+logger.remove()
+logger.add(sys.stderr, format="", diagnose=False, backtrace=False, colorize=False)
+logger.add(sys.stderr, format="", diagnose=True, backtrace=False, colorize=False)
+logger.add(sys.stderr, format="", diagnose=False, backtrace=True, colorize=False)
+logger.add(sys.stderr, format="", diagnose=True, backtrace=True, colorize=False)
+
+
+@logger.catch
+async def foo(a, b):
+    yield a / b
+
+
+f = foo(1, 0).asend(None)
+
+try:
+    f.send(None)
+except StopAsyncIteration:
+    pass

tests/test_exceptions_catch.py

@@ -407,9 +407,9 @@ def foo(x, y, z):
 def test_decorate_generator_with_error():
     @logger.catch
     def foo():
-        for i in range(3):
-            1 / (2 - i)
-            yield i
+        yield 0
+        yield 1
+        raise ValueError
 
     assert list(foo()) == [0, 1]
 

tests/test_exceptions_formatting.py

@@ -23,16 +23,21 @@ def normalize(exception):
 
         def fix_filepath(match):
             filepath = match.group(1)
+
+            # Pattern to check if the filepath contains ANSI escape codes.
             pattern = (
                 r'((?:\x1b\[[0-9]*m)+)([^"]+?)((?:\x1b\[[0-9]*m)+)([^"]+?)((?:\x1b\[[0-9]*m)+)'
             )
+
             match = re.match(pattern, filepath)
             start_directory = os.path.dirname(os.path.dirname(__file__))
             if match:
+                # Simplify the path while preserving the color highlighting of the file basename.
                 groups = list(match.groups())
                 groups[1] = os.path.relpath(os.path.abspath(groups[1]), start_directory) + "/"
                 relpath = "".join(groups)
             else:
+                # We can straightforwardly convert from absolute to relative path.
                 relpath = os.path.relpath(os.path.abspath(filepath), start_directory)
             return 'File "%s"' % relpath.replace("\\", "/")
 
@@ -241,6 +246,8 @@ def test_exception_others(filename):
     ("filename", "minimum_python_version"),
     [
         ("type_hints", (3, 6)),
+        ("exception_formatting_async_generator", (3, 6)),
+        ("decorate_async_generator", (3, 7)),
         ("positional_only_argument", (3, 8)),
         ("walrus_operator", (3, 8)),
         ("match_statement", (3, 10)),