[bug fix] add do_async to cover inlining multiple async calls in a do…

… expression

CHANGELOG.md

@@ -13,6 +13,8 @@ Possible log types:
 
 ## [Unreleased]
 
+- `[fixed]` Add `do_async()` to handle edge case in `do()` involving multiple inlined awaits (#149)
+
 
 ## [0.14.0] - 2023-11-10
 

README.rst

@@ -357,7 +357,7 @@ not the last.
         )
 
 NOTE: If you exclude the type annotation e.g. ``Result[float, int]``
-your type checker might be unable to infer the return type.
+your type checker might be unable to infer the type returned by ``do()``.
 To avoid an error, you might need to help it with the type hint.
 
 This is similar to Rust's  `m! macro <https://docs.rs/do-notation/latest/do_notation/>`_:
@@ -372,18 +372,32 @@ This is similar to Rust's  `m! macro <https://docs.rs/do-notation/latest/do_nota
     };
 
 
-You can also ``await`` Awaitables like async function calls. See example:
+You can use `do()` with awaited values as follows:
+
+
+.. sourcecode:: python
+
+    async def process_data(data) -> Result[float, int]:
+        res1 = await get_result_1(data)
+        res2 = await get_result_2(data)
+        return do(
+            Ok(len(x) + int(y) + 0.5)
+            for x in res1
+            for y in res2
+        )
+
+However, if you want to await something inside the expression,
+use `do_async`:
 
 .. sourcecode:: python
 
+
     async def process_data(data) -> Result[float, int]:
-        out: Result[float, int] = do(
+        return do_async(
             Ok(len(x) + int(y) + 0.5)
             for x in await get_result_1(data)
             for y in await get_result_2(data)
         )
-        return out
-
 
 Development
 ===========

src/result/__init__.py

@@ -9,6 +9,7 @@
     is_ok,
     is_err,
     do,
+    do_async,
 )
 
 __all__ = [
@@ -22,5 +23,6 @@
     "is_ok",
     "is_err",
     "do",
+    "do_async",
 ]
 __version__ = "0.15.0.dev0"

src/result/result.py

@@ -6,6 +6,7 @@
 from warnings import warn
 from typing import (
     Any,
+    AsyncGenerator,
     Awaitable,
     Callable,
     Final,
@@ -201,6 +202,7 @@ class DoException(Exception):
     to simulate `and_then()` in the Err case: namely, we don't call `op`,
     we just return `self` (the Err).
     """
+
     def __init__(self, err: Err[E]) -> None:
         self.err = err
 
@@ -537,3 +539,67 @@ def do(gen: Generator[Result[T, E], None, None]) -> Result[T, E]:
     except DoException as e:
         out: Err[E] = e.err  # type: ignore
         return out
+    except TypeError as te:
+        # Turn this into a more helpful error message.
+        # Python has strange rules involving turning generators involving `await`
+        # into async generators, so we want to make sure to help the user clearly.
+        if "'async_generator' object is not an iterator" in str(te):
+            raise TypeError(
+                "You have hit a case where `do()` does not support async calls.\n"
+                "Either use do_async(), or move your `await` expressions before `do()`\n."
+                "Example:\n"
+                "foo = await bar\ndo(...for foo_ok in foo)\n"
+            )
+        raise te
+
+
+async def do_async(
+    gen: Union[Generator[Result[T, E], None, None], AsyncGenerator[Result[T, E], None]]
+) -> Result[T, E]:
+    """Async version of do. Example:
+
+    >>> final_result: Result[float, int] = await do_async(
+        Ok(len(x) + int(y) + z)
+            for x in await get_async_result_1()
+            for y in await get_async_result_2()
+            for z in get_sync_result_3()
+        )
+
+    NOTE: Python makes generators async in a counter-intuitive way.
+    This is a regular generator:
+        async def foo(): ...
+        do(Ok(1) for x in await foo())
+
+    But this is an async generator:
+        async def foo(): ...
+        async def bar(): ...
+        do(
+            Ok(1)
+            for x in await foo()
+            for y in await bar()
+        )
+
+    We let users try to use regular `do()`, which works in some cases
+    of awaiting async values. If we hit a case like above, we raise
+    an exception telling the user to use `do_async()` instead.
+    See `do()`.
+
+    However, for better usability, it's better for `do_async()` to also accept
+    regular generators, as you get in the first case:
+
+    async def foo(): ...
+        do(Ok(1) for x in await foo())
+
+    Furthermore, neither mypy nor pyright can infer that the second case is
+    actually an async generator, so we cannot annotate `do_async()`
+    as accepting only an async generator. This is additional motivation
+    to accept either.
+    """
+    try:
+        if isinstance(gen, AsyncGenerator):
+            return await gen.__anext__()
+        else:
+            return next(gen)
+    except DoException as e:
+        out: Err[E] = e.err  # type: ignore
+        return out

tests/test_result_do.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from result import Err, Ok, Result, do
+from result import Err, Ok, Result, do, do_async
 
 
 def test_result_do_general() -> None:
@@ -24,25 +24,145 @@ def _get_output(is_suc1: bool, is_suc2: bool) -> Result[float, int]:
     assert _get_output(False, True) == Err(1)
     assert _get_output(False, False) == Err(1)
 
+    def _get_output_return_immediately(
+        is_suc1: bool, is_suc2: bool
+    ) -> Result[float, int]:
+        return do(
+            Ok(len(x) + int(y) + 0.5) for x in resx(is_suc1) for y in resy(is_suc2)
+        )
+
+    assert _get_output_return_immediately(True, True) == Ok(6.5)
+
 
 @pytest.mark.asyncio
-async def test_result_do_general_async() -> None:
-    async def get_resx(is_suc: bool) -> Result[str, int]:
+async def test_result_do_general_with_async_values() -> None:
+    # Asyncio works with regular `do()` as long as you await
+    # the async calls outside the `do()` expression.
+    # This causes the generator to be a regular (not async) generator.
+    async def aget_resx(is_suc: bool) -> Result[str, int]:
         return Ok("hello") if is_suc else Err(1)
 
-    async def get_resy(is_suc: bool) -> Result[bool, int]:
+    async def aget_resy(is_suc: bool) -> Result[bool, int]:
         return Ok(True) if is_suc else Err(2)
 
-    async def _get_output(is_suc1: bool, is_suc2: bool) -> Result[float, int]:
-        resx, resy = await get_resx(is_suc1), await get_resy(is_suc2)
+    async def _aget_output(is_suc1: bool, is_suc2: bool) -> Result[float, int]:
+        resx, resy = await aget_resx(is_suc1), await aget_resy(is_suc2)
         out: Result[float, int] = do(
-            Ok(len(x) + int(y) + 0.5)
-            for x in resx
-            for y in resy
+            Ok(len(x) + int(y) + 0.5) for x in resx for y in resy
+        )
+        return out
+
+    assert await _aget_output(True, True) == Ok(6.5)
+    assert await _aget_output(True, False) == Err(2)
+    assert await _aget_output(False, True) == Err(1)
+    assert await _aget_output(False, False) == Err(1)
+
+
+@pytest.mark.asyncio
+async def test_result_do_async_one_value() -> None:
+    """This is a strange case where Python creates a regular
+    (non async) generator despite an `await` inside the generator expression.
+    For convenience, although this works with regular `do()`, we want to support this
+    with `do_async()` as well."""
+
+    async def aget_resx(is_suc: bool) -> Result[str, int]:
+        return Ok("hello") if is_suc else Err(1)
+
+    def get_resz(is_suc: bool) -> Result[float, int]:
+        return Ok(0.5) if is_suc else Err(3)
+
+    assert await do_async(Ok(len(x)) for x in await aget_resx(True)) == Ok(5)
+    assert await do_async(Ok(len(x)) for x in await aget_resx(False)) == Err(1)
+
+    async def _aget_output(
+        is_suc1: bool, is_suc3: bool
+    ) -> Result[float, int]:
+        return await do_async(
+            Ok(len(x) + z) for x in await aget_resx(is_suc1) for z in get_resz(is_suc3)
+        )
+
+    assert await _aget_output(True, True) == Ok(5.5)
+    assert await _aget_output(True, False) == Err(3)
+    assert await _aget_output(False, True) == Err(1)
+    assert await _aget_output(False, False) == Err(1)
+
+
+@pytest.mark.asyncio
+async def test_result_do_async_general() -> None:
+    async def aget_resx(is_suc: bool) -> Result[str, int]:
+        return Ok("hello") if is_suc else Err(1)
+
+    async def aget_resy(is_suc: bool) -> Result[bool, int]:
+        return Ok(True) if is_suc else Err(2)
+
+    def get_resz(is_suc: bool) -> Result[float, int]:
+        return Ok(0.5) if is_suc else Err(3)
+
+    async def _aget_output(
+        is_suc1: bool, is_suc2: bool, is_suc3: bool
+    ) -> Result[float, int]:
+        out: Result[float, int] = await do_async(
+            Ok(len(x) + int(y) + z)
+            for x in await aget_resx(is_suc1)
+            for y in await aget_resy(is_suc2)
+            for z in get_resz(is_suc3)
         )
         return out
 
-    assert await _get_output(True, True) == Ok(6.5)
-    assert await _get_output(True, False) == Err(2)
-    assert await _get_output(False, True) == Err(1)
-    assert await _get_output(False, False) == Err(1)
+    assert await _aget_output(True, True, True) == Ok(6.5)
+    assert await _aget_output(True, False, True) == Err(2)
+    assert await _aget_output(False, True, True) == Err(1)
+    assert await _aget_output(False, False, True) == Err(1)
+
+    assert await _aget_output(True, True, False) == Err(3)
+    assert await _aget_output(True, False, False) == Err(2)
+    assert await _aget_output(False, True, False) == Err(1)
+    assert await _aget_output(False, False, False) == Err(1)
+
+    async def _aget_output_return_immediately(
+        is_suc1: bool, is_suc2: bool, is_suc3: bool
+    ) -> Result[float, int]:
+        return await do_async(
+            Ok(len(x) + int(y) + z)
+            for x in await aget_resx(is_suc1)
+            for y in await aget_resy(is_suc2)
+            for z in get_resz(is_suc3)
+        )
+
+    assert await _aget_output_return_immediately(True, True, True) == Ok(6.5)
+
+
+@pytest.mark.asyncio
+async def test_result_do_general_with_async_values_inline_error() -> None:
+    """
+    Due to subtle behavior, `do()` works in certain cases involving async
+    calls but not others. We surface a more helpful error to the user
+    in cases where it doesn't work indicating to use `do_async()` instead.
+    Contrast this with `test_result_do_general_with_async_values()`
+    in which using `do()` works with async functions as long as
+    their return values are resolved outside the `do()` expression.
+    """
+
+    async def aget_resx(is_suc: bool) -> Result[str, int]:
+        return Ok("hello") if is_suc else Err(1)
+
+    async def aget_resy(is_suc: bool) -> Result[bool, int]:
+        return Ok(True) if is_suc else Err(2)
+
+    def get_resz(is_suc: bool) -> Result[float, int]:
+        return Ok(0.5) if is_suc else Err(3)
+
+    with pytest.raises(TypeError) as excinfo:
+        do(
+            Ok(len(x) + int(y) + z)
+            for x in await aget_resx(True)
+            for y in await aget_resy(True)
+            for z in get_resz(True)
+        )
+
+    assert (
+        "You have hit a case where `do()` does not support async calls.\n"
+        "Either use do_async(), or move your `await` expressions before `do()`\n."
+        "Example:\n"
+        "foo = await bar\ndo(...for foo_ok in foo)\n"
+    ) in excinfo.value.args[0]