# Streams

class Stream:
    empty = 'empty'
    def __init__(self, first, compute_rest=lambda: Stream.empty):
        self.first = first
        self.cached_rest = None
        assert callable(compute_rest)
        self.compute_rest = compute_rest
    @property
    def rest(self):
        """Return the rest, computing it if necessary."""
        if self.compute_rest is not None:
            self.cached_rest = self.compute_rest()
            self.compute_rest = None
        return self.cached_rest
    def __repr__(self):
        rest = self.cached_rest if self.compute_rest is None else '<...>'
        return 'Stream({}, {})'.format(self.first, rest)

empty_stream = Stream.empty

def make_const_stream(x):
    """An infinite stream of X's."""
    return Stream(x, lambda: make_const_stream(x))

def make_integer_stream(first=1):
    """Return an infinite stream of increasing integers.

    >>> from operator import add
    >>> stream_to_list(make_integer_stream(1), 5)
    [1, 2, 3, 4, 5]
    """
    return Stream(first, lambda: make_integer_stream(first+1))

# Stream manipulation

def add_stream(s1, s2):
    if s1 is Stream.empty or s2 is Stream.empty:
        return Stream.empty
    else:
        return Stream(s1.first + s2.first, lambda: add_stream(s1.rest, s2.rest))

def map_stream(fn, s):
    """Return a stream of the values of fn applied to the elements of stream s.

    >>> s = make_integer_stream(3)
    >>> stream_to_list(truncate_stream(map_stream(lambda x: x*x, s), 4))
    [9, 16, 25, 36]
    """
    if s is Stream.empty:
        return s
    return Stream(fn(s.first), lambda: map_stream(fn, s.rest))

def filter_stream(fn, s):
    """Return a stream of the elements of s for which fn is true."""
    if s is Stream.empty:
        return s
    elif fn(s.first):
        return Stream(s.first, lambda: filter_stream(fn, s.rest))
    return filter_stream(fn, s.rest)

def stream_to_list(s, n=10):
    """A list containing the elements of stream S,
    up to a maximum of N."""
    r = []
    while n > 0 and s is not Stream.empty:
        r.append(s.first)
        s = s.rest
        n -= 1
    return r

def sieve(s):
    def not_divisible(x, y):
        return x % y != 0
    return Stream(s.first, 
        lambda: sieve(filter_stream(lambda x: not_divisible(x, s.first), s.rest)))

def fib_stream():
    return Stream(0, lambda: Stream(1, 
            lambda: add_stream(fib_stream(), fib_stream().rest)))

primes = sieve(make_integer_stream(2))
fib_sequence = fib_stream()