From c00ae7aca99a78c425dd983fcb68a371851b4037 Mon Sep 17 00:00:00 2001 From: ddzmanashvili Date: Thu, 30 May 2024 18:07:58 +0200 Subject: [PATCH] adapt `as_result` and `as_async_result` so they can accept generators --- CHANGELOG.md | 3 ++ src/result/result.py | 66 ++++++++++++++++++++++++++++++++------------ tests/test_result.py | 66 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 118 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 63f07aa..08ecb8e 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/src/result/result.py b/src/result/result.py index 06b78b9..72cce85 100644 --- a/src/result/result.py +++ b/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,28 +455,44 @@ 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)``. """ @@ -480,23 +500,35 @@ def as_async_result( 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 diff --git a/tests/test_result.py b/tests/test_result.py index f0e05c6..2af8fd0 100644 --- a/tests/test_result.py +++ b/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)