Python 3.12+ version of result.py from rustedpy/result

result.py

"""Result type for representing success and failure."""

# ruff: noqa: PLR6301
from __future__ import annotations

import functools
import inspect
from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Iterator
from typing import Any, Final, Literal, NoReturn

from typing_extensions import TypeIs


class Ok[T]:
    """A value that indicates success and which stores arbitrary data for the return value."""

    __match_args__ = ("ok_value",)
    __slots__ = ("_value",)

    def __iter__(self) -> Iterator[T]:
        """Return an iterator over the value."""
        yield self._value

    def __init__(self, value: T) -> None:
        """Initialize Ok with a value."""
        self._value = value

    def __repr__(self) -> str:
        """Return a string representation of the Ok value."""
        return f"Ok({self._value!r})"

    def __eq__(self, other: object) -> bool:
        """Return True if other is Ok and has the same value."""
        return isinstance(other, Ok) and self._value == other._value

    def __ne__(self, other: object) -> bool:
        """Return True if other is not equal to this Ok."""
        return not self == other

    def __hash__(self) -> int:
        """Return hash of the Ok value."""
        return hash((True, self._value))

    def is_ok(self) -> Literal[True]:
        """Return True since this is an Ok result."""
        return True

    def is_err(self) -> Literal[False]:
        """Return False since this is an Ok result."""
        return False

    def ok(self) -> T:
        """Return the value."""
        return self._value

    def err(self) -> None:
        """Return `None`."""
        return

    @property
    def ok_value(self) -> T:
        """Return the inner value."""
        return self._value

    def expect(self, _message: str) -> T:
        """Return the value."""
        return self._value

    def expect_err(self, message: str) -> NoReturn:
        """Raise an UnwrapError since this type is `Ok`."""
        raise UnwrapError(self, message)

    def unwrap(self) -> T:
        """Return the value."""
        return self._value

    def unwrap_err(self) -> NoReturn:
        """Raise an UnwrapError since this type is `Ok`."""
        raise UnwrapError(self, "Called `Result.unwrap_err()` on an `Ok` value")

    def unwrap_or(self, _default: object) -> T:
        """Return the value."""
        return self._value

    def unwrap_or_else(self, _op: object) -> T:
        """Return the value."""
        return self._value

    def unwrap_or_raise(self, _e: object) -> T:
        """Return the value."""
        return self._value

    def map[U](self, op: Callable[[T], U]) -> Ok[U]:
        """The contained result is `Ok`, so return `Ok` with original value mapped.

        Maps the original value to a new value using the passed in function.
        """
        return Ok(op(self._value))

    async def map_async[U](self, op: Callable[[T], Awaitable[U]]) -> Ok[U]:
        """The contained result is `Ok`, so return the result of `op`.

        Returns the result of calling `op` with the original value passed in.
        """
        return Ok(await op(self._value))

    def map_or[U](self, _default: object, op: Callable[[T], U]) -> U:
        """The contained result is `Ok`, so return the original value mapped to a new value.

        Apply the function using the passed in function.
        """
        return op(self._value)

    def map_or_else[U](self, _default_op: object, op: Callable[[T], U]) -> U:
        """The contained result is `Ok`, so return original value mapped to a new value.

        Apply the function using the passed in `op` function.
        """
        return op(self._value)

    def map_err(self, _op: object) -> Ok[T]:
        """The contained result is `Ok`, so return `Ok` with the original value."""
        return self

    def and_then[U, E](self, op: Callable[[T], Result[U, E]]) -> Result[U, E]:
        """The contained result is `Ok`, so return the result of `op` with the original value.

        Apply the provided operation to the value and return the result.
        """
        return op(self._value)

    async def and_then_async[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.

        Apply the provided operation to the value and return the result.
        """
        return await op(self._value)

    def or_else(self, _op: object) -> Ok[T]:
        """The contained result is `Ok`, so return `Ok` with the original value."""
        return self

    def inspect[E](self, op: Callable[[T], Any]) -> Result[T, E]:
        """Calls a function with the contained value if `Ok`. Returns the original result."""
        op(self._value)
        return self

    def inspect_err[E](self, _op: Callable[[E], Any]) -> Result[T, E]:
        """Calls a function with the contained value if `Err`. Returns the original result."""
        return self


class DoError[E](Exception):
    """This is used to signal to `do()` that the result is an `Err`.

    This 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:
        """Initialize DoError with an Err value."""
        self.err = err


class Err[E]:
    """A value that signifies failure and which stores arbitrary data for the error."""

    __match_args__ = ("err_value",)
    __slots__ = ("_value",)

    def __iter__(self) -> Iterator[NoReturn]:
        """Return an iterator that raises DoException when advanced."""

        def _iter() -> Iterator[NoReturn]:
            # Exception will be raised when the iterator is advanced, not when it's created
            raise DoError(self)
            # pylint: disable=unreachable
            yield  # This yield will never be reached, but is necessary to create a generator

        return _iter()

    def __init__(self, value: E) -> None:
        """Initialize Err with an error value."""
        self._value = value

    def __repr__(self) -> str:
        """Return a string representation of the Err value."""
        return f"Err({self._value!r})"

    def __eq__(self, other: object) -> bool:
        """Return True if other is Err and has the same value."""
        return isinstance(other, Err) and self._value == other._value

    def __ne__(self, other: object) -> bool:
        """Return True if other is not equal to this Err."""
        return not self == other

    def __hash__(self) -> int:
        """Return hash of the Err value."""
        return hash((False, self._value))

    def is_ok(self) -> Literal[False]:
        """Return False since this is an Err result."""
        return False

    def is_err(self) -> Literal[True]:
        """Return True since this is an Err result."""
        return True

    def ok(self) -> None:
        """Return `None`."""
        return

    def err(self) -> E:
        """Return the error."""
        return self._value

    @property
    def err_value(self) -> E:
        """Return the inner value."""
        return self._value

    def expect(self, message: str) -> NoReturn:
        """Raises an `UnwrapError`."""
        exc = UnwrapError(
            self,
            f"{message}: {self._value!r}",
        )
        if isinstance(self._value, BaseException):
            raise exc from self._value
        raise exc

    def expect_err(self, _message: str) -> E:
        """Return the inner value."""
        return self._value

    def unwrap(self) -> NoReturn:
        """Raises an `UnwrapError`."""
        exc = UnwrapError(
            self,
            f"Called `Result.unwrap()` on an `Err` value: {self._value!r}",
        )
        if isinstance(self._value, BaseException):
            raise exc from self._value
        raise exc

    def unwrap_err(self) -> E:
        """Return the inner value."""
        return self._value

    def unwrap_or[U](self, default: U) -> U:
        """Return `default`."""
        return default

    def unwrap_or_else[T](self, op: Callable[[E], T]) -> T:
        """The contained result is ``Err``, so return the result of applying ``op`` to the error value.

        Apply the provided operation to the error value and return the result.
        """
        return op(self._value)

    def unwrap_or_raise[TBE: BaseException](self, e: type[TBE]) -> NoReturn:
        """The contained result is ``Err``, so raise the exception with the value."""
        raise e(self._value)

    def map(self, _op: object) -> Err[E]:
        """Return `Err` with the same value."""
        return self

    async def map_async(self, _op: object) -> Err[E]:
        """The contained result is `Err`, so return `Err` with the original value.

        Return the original value passed in.
        """
        return self

    def map_or[U](self, default: U, _op: object) -> U:
        """Return the default value."""
        return default

    def map_or_else[U](self, default_op: Callable[[], U], _op: object) -> U:
        """Return the result of the default operation."""
        return default_op()

    def map_err[F](self, op: Callable[[E], F]) -> Err[F]:
        """The contained result is `Err`, so return `Err` with original error mapped to a new value.

        Apply the passed in function to map the error to a new value.
        """
        return Err(op(self._value))

    def and_then(self, _op: object) -> Err[E]:
        """The contained result is `Err`, so return `Err` with the original value."""
        return self

    async def and_then_async(self, _op: object) -> Err[E]:
        """The contained result is `Err`, so return `Err` with the original value."""
        return self

    def or_else[T, F](self, op: Callable[[E], Result[T, F]]) -> Result[T, F]:
        """The contained result is `Err`, so return the result of `op` with the original value.

        Apply the provided operation to the error value and return the result.
        """
        return op(self._value)

    def inspect[T](self, _op: Callable[[T], Any]) -> Result[T, E]:
        """Calls a function with the contained value if `Ok`. Returns the original result."""
        return self

    def inspect_err[T](self, op: Callable[[E], Any]) -> Result[T, E]:
        """Calls a function with the contained value if `Err`. Returns the original result."""
        op(self._value)
        return self


# define Result as a generic type alias for use
# in type annotations
type Result[_T, _E] = Ok[_T] | Err[_E]
"""
A simple `Result` type inspired by Rust.
Not all methods (https://doc.rust-lang.org/std/result/enum.Result.html)
have been implemented, only the ones that make sense in the Python context.
"""

OkErr: Final = (Ok, Err)
"""
A type to use in `isinstance` checks.
This is purely for convenience sake, as you could also just write `isinstance(res, (Ok, Err))
"""


class UnwrapError(Exception):
    """Exception raised from ``.unwrap_<...>`` and ``.expect_<...>`` calls.

    The original ``Result`` can be accessed via the ``.result`` attribute, but
    this is not intended for regular use, as type information is lost:
    ``UnwrapError`` doesn't know about both ``T`` and ``E``, since it's raised
    from ``Ok()`` or ``Err()`` which only knows about either ``T`` or ``E``,
    not both.
    """

    _result: Result[object, object]

    def __init__(self, result: Result[object, object], message: str) -> None:
        """Initialize UnwrapError with a result and message."""
        self._result = result
        super().__init__(message)

    @property
    def result(self) -> Result[Any, Any]:
        """Returns the original result."""
        return self._result


def as_result[R, **P, TBE: BaseException](
    *exceptions: type[TBE],
) -> Callable[[Callable[P, R]], Callable[P, Result[R, TBE]]]:
    """Make a decorator to turn a function 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")

    def decorator(f: Callable[P, R]) -> Callable[P, Result[R, TBE]]:
        """Decorator to turn a function into one that returns a ``Result``."""

        @functools.wraps(f)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> Result[R, TBE]:
            try:
                return Ok(f(*args, **kwargs))
            except exceptions as exc:
                return Err(exc)

        return wrapper

    return decorator


def as_async_result[R, **P, TBE: BaseException](
    *exceptions: type[TBE],
) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[Result[R, TBE]]]]:
    """Make a decorator to turn an async function 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")

    def decorator(f: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[Result[R, TBE]]]:
        """Decorator to turn a function into one that returns a ``Result``."""

        @functools.wraps(f)
        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Result[R, TBE]:
            try:
                return Ok(await f(*args, **kwargs))
            except exceptions as exc:
                return Err(exc)

        return async_wrapper

    return decorator


def is_ok[T, E](result: Result[T, E]) -> TypeIs[Ok[T]]:
    """A type guard to check if a result is an Ok.

    Usage:

    ``` python
    r: Result[int, str] = get_a_result()
    if is_ok(r):
        r  # r is of type Ok[int]
    elif is_err(r):
        r  # r is of type Err[str]
    ```

    """
    return result.is_ok()


def is_err[T, E](result: Result[T, E]) -> TypeIs[Err[E]]:
    """A type guard to check if a result is an Err.

    Usage:

    ``` python
    r: Result[int, str] = get_a_result()
    if is_ok(r):
        r  # r is of type Ok[int]
    elif is_err(r):
        r  # r is of type Err[str]
    ```

    """
    return result.is_err()


def do[T, E](gen: Generator[Result[T, E], None, None]) -> Result[T, E]:
    """Do notation for Result (syntactic sugar for sequence of `and_then()` calls).

    Usage:

    ```rust
    // 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)
    };
    ```

    ```rust
    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 DoError 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(
                "Got async_generator but expected generator.See the section on do notation in the README."
            ) from te
        raise


async def do_async[T, E](
    gen: Generator[Result[T, E], None, None] | AsyncGenerator[Result[T, E], None],
) -> Result[T, E]:
    """Async version of do.

    ```python
    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.

    ```python
    # This is a regular generator:
        async def foo(): ...
        do(Ok(1) for x in await foo())
    ```

    ```python
    # 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:

    ```python
    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 anext(gen)
        return next(gen)
    except DoError as e:
        out: Err[E] = e.err  # type: ignore
        return out