Repository: Delgan/loguru (python)

Issue:
Make "logger.catch()" compatible with asynchronous generators
Eg in:

```python
import asyncio
from loguru import logger


@logger.catch
async def foo():
    yield 1
    yield 2
    raise ValueError


async def main():
    async for i in foo():
        print(i)

    print("Done")


asyncio.run(main())
```

Pull request description:
Add support for async generators to "@logger.catch"
Fix #1302.

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).

Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 5e0e34d..1f575b8 100644
--- a/CHANGELOG.rst
+++ b/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)
 =====================
diff --git a/loguru/_logger.py b/loguru/_logger.py
index 911dd60..5c2ba54 100644
--- a/loguru/_logger.py
+++ b/loguru/_logger.py
@@ -114,10 +114,16 @@ from ._recattrs import RecordException, RecordFile, RecordLevel, RecordProcess,
 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 @@ class Logger:
                             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):
diff --git a/pyproject.toml b/pyproject.toml
index 269b914..c689c01 100644
--- a/pyproject.toml
+++ b/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.


Test patch:
diff --git a/tests/exceptions/output/modern/decorate_async_generator.txt b/tests/exceptions/output/modern/decorate_async_generator.txt
new file mode 100644
index 0000000..a965a70
--- /dev/null
+++ b/tests/exceptions/output/modern/decorate_async_generator.txt
@@ -0,0 +1,1 @@
+Done
diff --git a/tests/exceptions/output/modern/exception_formatting_async_generator.txt b/tests/exceptions/output/modern/exception_formatting_async_generator.txt
new file mode 100644
index 0000000..c5fdfab
--- /dev/null
+++ b/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
diff --git a/tests/exceptions/source/modern/decorate_async_generator.py b/tests/exceptions/source/modern/decorate_async_generator.py
new file mode 100644
index 0000000..b1dffed
--- /dev/null
+++ b/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")
diff --git a/tests/exceptions/source/modern/exception_formatting_async_generator.py b/tests/exceptions/source/modern/exception_formatting_async_generator.py
new file mode 100644
index 0000000..120bc4c
--- /dev/null
+++ b/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
diff --git a/tests/test_exceptions_catch.py b/tests/test_exceptions_catch.py
index 1e7aea9..71ebc29 100644
--- a/tests/test_exceptions_catch.py
+++ b/tests/test_exceptions_catch.py
@@ -407,9 +407,9 @@ def test_decorate_generator(writer):
 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]
 
diff --git a/tests/test_exceptions_formatting.py b/tests/test_exceptions_formatting.py
index 47e05cf..58b6646 100644
--- a/tests/test_exceptions_formatting.py
+++ b/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)),