Add generator-based do notation (#150)

.gitignore

@@ -11,3 +11,5 @@ build/
 .mypy_cache/
 venv/
 /.tox/
+.vscode
+pyrightconfig.json

CHANGELOG.md

@@ -18,6 +18,8 @@ Possible log types:
 
 - `[added]` `is_ok` and `is_err` type guard functions as alternatives to `isinstance` checks (#69)
 - `[added]` Add `and_then_async` for async functions (#148)
+- `[added]` Support do notation (#149)
+
 
 ## [0.13.1] - 2023-07-19
 

README.rst

@@ -340,6 +340,61 @@ unconventional syntax (without the usual ``@``):
     if isinstance(res, Ok):
         print(res.value)
 
+
+Do notation: syntactic sugar for a sequence of ``and_then()`` calls.
+Much like the equivalent in Rust or Haskell, but with different syntax.
+Instead of ``x <- Ok(1)`` we write ``for x in Ok(1)``.
+Since the syntax is generator-based, the final result must be the first line,
+not the last.
+
+.. sourcecode:: python
+
+
+    >>> final_result: Result[float, int] = do(
+            Ok(len(x) + int(y) + 0.5)
+            for x in Ok("hello")
+            for y in Ok(True)
+        )
+
+NOTE: If you exclude the type annotation e.g. ``Result[float, int]``
+your type checker might be unable to infer the return type.
+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/>`_:
+
+.. sourcecode:: rust
+
+    use do_notation::m;
+    let r = m! {
+        x <- Some("Hello, world!");
+        y <- Some(3);
+        Some(x.len() * y)
+    };
+
+
+You can also ``await`` Awaitables like async function calls. See example:
+
+.. sourcecode:: python
+
+    async def process_data(data) -> Result[float, int]:
+        out: Result[float, int] = do(
+            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
+===========
+
+* Set up: ``pip install -e .``
+
+* Test your changes: ``flake8 src tests; mypy; pytest``
+
+* Remember to test in Python 3.8 - 3.11.
+
+
 FAQ
 ===
 

src/result/__init__.py

@@ -8,6 +8,7 @@
     as_result,
     is_ok,
     is_err,
+    do,
 )
 
 __all__ = [
@@ -20,5 +21,6 @@
     "as_result",
     "is_ok",
     "is_err",
+    "do",
 ]
 __version__ = "0.15.0.dev0"

src/result/result.py

@@ -9,7 +9,9 @@
     Awaitable,
     Callable,
     Final,
+    Generator,
     Generic,
+    Iterator,
     Literal,
     NoReturn,
     Type,
@@ -40,6 +42,9 @@ class Ok(Generic[T]):
     __match_args__ = ("ok_value",)
     __slots__ = ("_value",)
 
+    def __iter__(self) -> Iterator[T]:
+        yield self._value
+
     def __init__(self, value: T) -> None:
         self._value = value
 
@@ -173,8 +178,8 @@ def and_then(self, op: Callable[[T], Result[U, E]]) -> Result[U, E]:
         return op(self._value)
 
     async def and_then_async(
-            self,
-            op: Callable[[T], Awaitable[Result[U, E]]]) -> Result[U, E]:
+        self, op: Callable[[T], Awaitable[Result[U, E]]]
+    ) -> Result[U, E]:
         """
         The contained result is `Ok`, so return the result of `op` with the
         original value passed in
@@ -188,6 +193,18 @@ def or_else(self, op: object) -> Ok[T]:
         return self
 
 
+class DoException(Exception):
+    """
+    This is used to signal to `do()` that the result is an `Err`,
+    which short-circuits the generator and returns that Err.
+    Using this exception for control flow in `do()` allows us
+    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
+
+
 class Err(Generic[E]):
     """
     A value that signifies failure and which stores arbitrary data for the error.
@@ -196,6 +213,14 @@ class Err(Generic[E]):
     __match_args__ = ("err_value",)
     __slots__ = ("_value",)
 
+    def __iter__(self) -> Iterator[NoReturn]:
+        def _iter() -> Iterator[NoReturn]:
+            # Exception will be raised when the iterator is advanced, not when it's created
+            raise DoException(self)
+            yield  # This yield will never be reached, but is necessary to create a generator
+
+        return _iter()
+
     def __init__(self, value: E) -> None:
         self._value = value
 
@@ -482,3 +507,33 @@ def is_err(result: Result[T, E]) -> TypeGuard[Err[E]]:
     >>>     r   # r is of type Err[str]
     """
     return result.is_err()
+
+
+def do(gen: Generator[Result[T, E], None, None]) -> Result[T, E]:
+    """Do notation for Result (syntactic sugar for sequence of `and_then()` calls).
+
+
+    Usage:
+    >>> # This is similar to
+    use do_notation::m;
+    let final_result = m! {
+        x <- Ok("hello");
+        y <- Ok(True);
+        Ok(len(x) + int(y) + 0.5)
+    };
+
+    >>> final_result: Result[float, int] = do(
+            Ok(len(x) + int(y) + 0.5)
+            for x in Ok("hello")
+            for y in Ok(True)
+        )
+
+    NOTE: If you exclude the type annotation e.g. `Result[float, int]`
+    your type checker might be unable to infer the return type.
+    To avoid an error, you might need to help it with the type hint.
+    """
+    try:
+        return next(gen)
+    except DoException as e:
+        out: Err[E] = e.err  # type: ignore
+        return out

tests/test_result_do.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+
+import pytest
+
+from result import Err, Ok, Result, do
+
+
+def test_result_do_general() -> None:
+    def resx(is_suc: bool) -> Result[str, int]:
+        return Ok("hello") if is_suc else Err(1)
+
+    def resy(is_suc: bool) -> Result[bool, int]:
+        return Ok(True) if is_suc else Err(2)
+
+    def _get_output(is_suc1: bool, is_suc2: bool) -> Result[float, int]:
+        out: Result[float, int] = do(
+            Ok(len(x) + int(y) + 0.5) for x in resx(is_suc1) for y in resy(is_suc2)
+        )
+        return out
+
+    assert _get_output(True, True) == Ok(6.5)
+    assert _get_output(True, False) == Err(2)
+    assert _get_output(False, True) == Err(1)
+    assert _get_output(False, False) == Err(1)
+
+
+@pytest.mark.asyncio
+async def test_result_do_general_async() -> None:
+    async def get_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]:
+        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)
+        out: Result[float, int] = do(
+            Ok(len(x) + int(y) + 0.5)
+            for x in resx
+            for y in resy
+        )
+        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)