adapt as_result and as_async_result so they can accept generators

CHANGELOG.md

@@ -13,6 +13,9 @@ Possible log types:
 
 ## [Unreleased]
 
+- `[changed]` changed `as_result` so it can work with both generators and functions
+- `[changed]` changed `as_async_result` so it can work with both async generators and async functions
+
 ## [0.16.1] - 2024-02-29
 
 - `[fixed]` PyPI not showing description (#176)

src/result/result.py

@@ -437,10 +437,14 @@ def result(self) -> Result[Any, Any]:
 
 
 def as_result(
-    *exceptions: Type[TBE],
-) -> Callable[[Callable[P, R]], Callable[P, Result[R, TBE]]]:
+    *exceptions: Type[TBE]
+) -> Callable[
+    [Callable[P, R] | Callable[P, Generator[R, Any, Any]]],
+    Callable[P, Result[R, TBE]] | Callable[P, Generator[Result[R, TBE], Any, Any]]
+]:
     """
-    Make a decorator to turn a function into one that returns a ``Result``.
+    Make a decorator to turn a function or generator into one that returns a ``Result``
+    or Generator yielding ``Result``.
 
     Regular return values are turned into ``Ok(return_value)``. Raised
     exceptions of the specified exception type(s) are turned into ``Err(exc)``.
@@ -451,52 +455,80 @@ def as_result(
     ):
         raise TypeError("as_result() requires one or more exception types")
 
-    def decorator(f: Callable[P, R]) -> Callable[P, Result[R, TBE]]:
+    def decorator(
+        f: Callable[P, R] | Callable[P, Generator[R, Any, Any]]
+    ) -> Callable[P, Result[R, TBE]] | Callable[P, Generator[Result[R, TBE], Any, Any]]:
         """
-        Decorator to turn a function into one that returns a ``Result``.
+        Decorator to turn a function or generator into one that returns a ``Result``.
         """
 
         @functools.wraps(f)
-        def wrapper(*args: P.args, **kwargs: P.kwargs) -> Result[R, TBE]:
+        def wrapper_function(*args: P.args, **kwargs: P.kwargs) -> Result[R, TBE]:
+            assert inspect.isfunction(f)
             try:
                 return Ok(f(*args, **kwargs))
             except exceptions as exc:
                 return Err(exc)
 
-        return wrapper
+        @functools.wraps(f)
+        def wrapper_generator(*args: P.args, **kwargs: P.kwargs) -> Generator[Result[R, TBE], Any, Any]:
+            assert inspect.isgeneratorfunction(f)
+            try:
+                for value in f(*args, **kwargs):
+                    yield Ok(value)
+            except exceptions as exc:
+                yield Err(exc)
+
+        if inspect.isgeneratorfunction(f):
+            return wrapper_generator
+        return wrapper_function
 
     return decorator
 
 
-def as_async_result(
-    *exceptions: Type[TBE],
-) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[Result[R, TBE]]]]:
+def as_async_result(*exceptions: Type[TBE]) -> Callable[
+    [Callable[P, Awaitable[R]] | Callable[P, AsyncGenerator[R, Any]]],
+    Callable[P, Awaitable[Result[R, TBE]]] | Callable[P, AsyncGenerator[Result[R, TBE], Any,]]
+]:
     """
-    Make a decorator to turn an async function into one that returns a ``Result``.
+    Make a decorator to turn an async function or async generator into one that returns a ``Result``.
+
     Regular return values are turned into ``Ok(return_value)``. Raised
     exceptions of the specified exception type(s) are turned into ``Err(exc)``.
     """
     if not exceptions or not all(
         inspect.isclass(exception) and issubclass(exception, BaseException)
         for exception in exceptions
     ):
-        raise TypeError("as_result() requires one or more exception types")
+        raise TypeError("as_async_result() requires one or more exception types")
 
     def decorator(
-        f: Callable[P, Awaitable[R]]
-    ) -> Callable[P, Awaitable[Result[R, TBE]]]:
+        f: Callable[P, Awaitable[R]] | Callable[P, AsyncGenerator[R, Any]]
+    ) -> Callable[P, Awaitable[Result[R, TBE]]] | Callable[P, AsyncGenerator[Result[R, TBE], Any]]:
         """
-        Decorator to turn a function into one that returns a ``Result``.
+        Decorator to turn an async function into one that returns a ``Result``.
         """
 
         @functools.wraps(f)
-        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Result[R, TBE]:
+        async def async_wrapper_function(*args: P.args, **kwargs: P.kwargs) -> Result[R, TBE]:
+            assert inspect.iscoroutinefunction(f)
             try:
                 return Ok(await f(*args, **kwargs))
             except exceptions as exc:
                 return Err(exc)
 
-        return async_wrapper
+        @functools.wraps(f)
+        async def async_wrapper_generator(*args: P.args, **kwargs: P.kwargs) -> AsyncGenerator[Result[R, TBE], Any]:
+            assert inspect.isasyncgenfunction(f)
+            try:
+                async for value in f(*args, **kwargs):
+                    yield Ok(value)
+            except exceptions as exc:
+                yield Err(exc)
+
+        if inspect.isasyncgenfunction(f):
+            return async_wrapper_generator
+        return async_wrapper_function
 
     return decorator
 

tests/test_result.py

@@ -306,6 +306,38 @@ def bad(value: int) -> int:
     assert isinstance(bad_result.unwrap_err(), ValueError)
 
 
+def test_as_result_with_generator() -> None:
+    """
+    ``as_result()`` works with generators.
+    """
+
+    @as_result(ValueError)
+    def random_generator():
+        yield 1
+        yield 2
+        yield 3
+
+    for i in random_generator():
+        assert i in (Ok(1), Ok(2), Ok(3))
+
+
+def test_as_result_with_generator_and_exception() -> None:
+    """
+    ``as_result()`` works with generators that raise exceptions.
+    """
+
+    @as_result(ValueError)
+    def random_generator():
+        yield 1
+        yield 2
+        raise ValueError
+
+    result = [i for i in random_generator()]
+    assert result[:2] == [Ok(1), Ok(2)]
+    assert result[-1].is_err()
+    assert isinstance(result[-1].unwrap_err(), ValueError)
+
+
 def test_as_result_other_exception() -> None:
     """
     ``as_result()`` only catches the specified exceptions.
@@ -375,6 +407,40 @@ async def bad(value: int) -> int:
     assert isinstance(bad_result.unwrap_err(), ValueError)
 
 
+@pytest.mark.asyncio
+async def test_as_async_result_with_generator() -> None:
+    """
+    ``as_result()`` works with async generators.
+    """
+
+    @as_async_result(ValueError)
+    async def random_generator():
+        yield 1
+        yield 2
+        yield 3
+
+    async for i in random_generator():
+        assert i in (Ok(1), Ok(2), Ok(3))
+
+
+@pytest.mark.asyncio
+async def test_as_async_result_with_generator_and_exception() -> None:
+    """
+    ``as_result()`` works with async generators that raise exceptions.
+    """
+
+    @as_async_result(ValueError)
+    async def random_generator():
+        yield 1
+        yield 2
+        raise ValueError
+
+    result = [i async for i in random_generator()]
+    assert result[:2] == [Ok(1), Ok(2)]
+    assert result[-1].is_err()
+    assert isinstance(result[-1].unwrap_err(), ValueError)
+
+
 def sq(i: int) -> Result[int, int]:
     return Ok(i * i)