tco python with exceptions

tco.py

from typing import ParamSpec, TypeVar, Callable, reveal_type


T = TypeVar("T")
P = ParamSpec("P")


def tco(f: Callable[P, T]) -> Callable[P, T]:
    unwind = False

    class Unwind(Exception):
        args: P.args
        kwargs: P.kwargs

        def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
            self.args = args
            self.kwargs = kwargs

    def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
        nonlocal unwind
        if unwind:
            raise Unwind(*args, **kwargs)

        unwind = True
        while True:
            try:
                result = f(*args, **kwargs)
            except Unwind as s:
                args = s.args
                kwargs = s.kwargs
                continue
            else:
                unwind = False
                return result

    return _inner


@tco
def recurse(x: int) -> str:
    if x <= 0:
        return "ok"
    return recurse(x - 1)


@tco
def sum(x: int, acc: int) -> int:
    if x <= 0:
        return acc
    return sum(x - 1, acc + x)


print(recurse(1001))  # ok
print(recurse(1001))  # ok
print(sum(1001, 0))