Understanding the Haskell IO monad by implementing it in Python

pymonad.py

# Based on https://youtu.be/fCoQb-zqYDI

# This passes `mypy --strict`, which is cool.

from __future__ import annotations

from typing import NamedTuple, Tuple, Callable, TypeVar

A = TypeVar('A')
B = TypeVar('B')


class World(NamedTuple):
    pass


def print_str(s: str, w: World) -> World:
    print(s)
    return w


def read_str(w: World) -> Tuple[str, World]:
    s = input()
    return s, w


WorldT = Callable[[World], Tuple[A, World]]


def bind(wt: WorldT[A], f: Callable[[A], WorldT[B]]) -> WorldT[B]:
    def wt_r(w: World) -> Tuple[B, World]:
        t, w2 = wt(w)
        wt2 = f(t)
        s, w3 = wt2(w2)
        return s, w3
    return wt_r


def read_str_t(w: World) -> Tuple[str, World]:
    s, w = read_str(w)
    return s, w


_verify_read_str_t_type: WorldT[str] = read_str_t


def print_str_t(s: str) -> WorldT[None]:
    def wt(w: World) -> Tuple[None, World]:
        w2 = print_str(s, w)
        return None, w2
    return wt


def what_is_your_pure_name(w1: World) -> World:
    w2 = print_str("What is your name?", w1)
    name, w3 = read_str(w2)
    w4 = print_str(f"Hello, {name}", w3)
    return w4


what_is_your_pure_name_t: WorldT[None] = bind(
    print_str_t("What is your name?"),
    lambda _: bind(read_str_t, lambda name: print_str_t(f"Hi, {name}")))


conditional_t: WorldT[None] = bind(
    print_str_t("What is your name?"),
    lambda _: bind(
        read_str_t, lambda name: bind(
            print_str_t("That's a nice name!") if name == 'noam' else print_str_t("That's a name."),
            lambda _: print_str_t("Ok, I'm done.")
        )
    )
)