Representing a stream in Python. The idea is came from the MIT 6.001 Lec6B.

stream.py

#!/usr/bin/env python3

def cons_stream(head, tail):
    return (head, delay(tail))


def head(s):
    return s[0]


def tail(s):
    return force(s[1])


def memo_proc(proc):
    already_run = False
    result = None

    def inner():
        nonlocal already_run
        nonlocal result
        if not already_run:
            result = proc()
            already_run = True
        return result
    return inner


def delay(proc):
    return memo_proc(proc)


def force(proc):
    return proc()


def add_streams(s1, s2):
    if len(s1) == 0:
        return s2
    elif len(s2) == 0:
        return s1

    return cons_stream(
        head(s1) + head(s2),
        lambda: add_streams(tail(s1), tail(s2))
    )


def print_stream(s):
    if len(s) == 0:
        print("Done")
    else:
        print(f"{head(s)} ", end="")
        print_stream(tail(s))


def nth_elem(n, s):
    while n > 0:
        s = tail(s)
        n -= 1
    return head(s)


def map_stream(p, s):
    return cons_stream(
        p(head(s)),
        lambda: map_stream(p, tail(s))
    )


def filter_stream(p, s):
    if len(s) == 0:
        return ()
    elif p(head(s)):
        return cons_stream(
            head(s),
            lambda: filter_stream(p, tail(s))
        )
    else:
        return filter_stream(p, tail(s))


if __name__ == "__main__":
    ones = cons_stream(1, lambda: ones)
    integers = cons_stream(1, lambda: add_streams(ones, integers))
    squares = map_stream(lambda n: n ** 2, integers)
    even_squares = filter_stream(lambda n: n % 2 == 0, squares)
    print(nth_elem(10, even_squares))