Add flatmap to Result

Author: ilotoki0804

src/fieldenum/enums.py

@@ -27,7 +27,7 @@
 from . import Unit, Variant, fieldenum, unreachable
 from .exceptions import IncompatibleBoundError, UnwrapFailedError
 
-__all__ = ["Option", "BoundResult", "Message", "Some", "Success", "Failed"]
+__all__ = ["Option", "BoundResult", "Message", "Some", "Success", "Failed", "Result", "Ok", "Err"]
 
 _MISSING = object()
 type _ExceptionTypes = type[BaseException] | tuple[type[BaseException], ... ] | UnionType
@@ -347,15 +347,15 @@ def dump(self) -> tuple[E]: ...
     def __bool__(self) -> bool:
         return isinstance(self, Result.Ok)
 
-    @overload
-    def unwrap(self) -> R: ...
-
     @overload
     def unwrap(self, default: R) -> R: ...
 
     @overload
     def unwrap[T](self, default: T) -> R | T: ...
 
+    @overload
+    def unwrap(self) -> R: ...
+
     def unwrap(self, default=_MISSING):
         match self:
             case Result.Ok(value):
@@ -387,7 +387,7 @@ def as_option(self) -> Option[R]:
     def exit(self, error_code: str | int | None = 1) -> NoReturn:
         sys.exit(0 if self else error_code)
 
-    def map[NewReturn](self, func: Callable[[R], NewReturn], bound: _ExceptionTypes, /) -> Result[NewReturn, E]:
+    def map[NewReturn](self, func: Callable[[R], NewReturn], /, bound: _ExceptionTypes) -> Result[NewReturn, E]:
         match self:
             case Result.Ok(ok):
                 try:
@@ -398,11 +398,32 @@ def map[NewReturn](self, func: Callable[[R], NewReturn], bound: _ExceptionTypes,
                     else:
                         raise
 
-            case Result.Err(error) as failed:
-                if TYPE_CHECKING:
-                    return Result.Err[NewReturn, E](error)
+            case Result.Err() as err:
+                return err  # type: ignore
+
+            case other:
+                unreachable(other)
+
+    def flatmap[NewResult: Result](self, func: Callable[[R], NewResult], /, bound: _ExceptionTypes) -> NewResult:
+        match self:
+            case Result.Ok(value):
+                try:
+                    result = func(value)
+                except BaseException as exc:
+                    if isinstance(exc, bound):
+                        return Result.Err(exc)  # type: ignore
+                    else:
+                        raise
+
+                if isinstance(result, Result):
+                    return result
                 else:
-                    return failed
+                    raise TypeError(
+                        f"Expect Result but received {type(result).__name__!r}"
+                    )
+
+            case Result.Err() as err:
+                return err  # type: ignore
 
             case other:
                 unreachable(other)
@@ -663,5 +684,7 @@ def dump(self) -> tuple[()]: ...
 
 
 Some = Option.Some
+Ok = Result.Ok
+Err = Result.Err
 Success = BoundResult.Success
 Failed = BoundResult.Failed

tests/test_enums.py

@@ -1,11 +1,52 @@
+from collections.abc import Callable
 from typing import Any, assert_type
 import pytest
 from fieldenum import *
 from fieldenum.enums import *
 from fieldenum.exceptions import IncompatibleBoundError, UnwrapFailedError
 
 
-def test_option_spread_map():
+def test_result_maps():
+    result = Result.Ok("hello")
+    err = result.map(int, ValueError)
+    assert type(err.error) == ValueError  # type: ignore
+    upper = result.map(str.upper, Exception)
+    assert upper.unwrap() == "HELLO"
+    with pytest.raises(ValueError):
+        result.map(int, ArithmeticError)
+
+    def returns_result[T, R, E: BaseException](func: Callable[[T], R], bound: type[E]) -> Callable[[T], Result[R, E]]:
+        def inner(value: T):
+            try:
+                return Ok(func(value))
+            except bound as exc:
+                return Err(exc)
+        return inner
+
+    # flatmap with returning result
+    result = Result.Ok("hello")
+    err = result.flatmap(returns_result(int, ValueError), ValueError)
+    assert type(err.error) == ValueError  # type: ignore
+    upper = result.flatmap(returns_result(str.upper, ValueError), Exception)
+    assert upper.unwrap() == "HELLO"
+    err = Exception()
+    assert Result.Err(err).flatmap(returns_result(int, ValueError), ValueError).error is err  # type: ignore
+    with pytest.raises(ValueError):
+        result.flatmap(returns_result(int, ArithmeticError), ArithmeticError)  # type: ignore
+
+    # flatmap with raising
+    err = result.flatmap(returns_result(int, ArithmeticError), ValueError)
+    assert type(err.error) == ValueError  # type: ignore
+    upper = result.flatmap(returns_result(str.upper, ArithmeticError), Exception)
+    assert upper.unwrap() == "HELLO"
+    assert Result.Err(err).flatmap(returns_result(int, ArithmeticError), ValueError).error is err  # type: ignore
+    with pytest.raises(ValueError):
+        result.flatmap(returns_result(int, ArithmeticError), ArithmeticError)  # type: ignore
+    with pytest.raises(TypeError):
+        result.flatmap(str.upper, ArithmeticError)  # type: ignore
+
+
+def test_option_flatmap():
     opt = Option.new("123").map(int).flatmap(Option.new)
     assert_type(opt, Option[int])
     assert opt == Option.Some(123)