PrimeStream

#!/usr/bin/env python
from collections import defaultdict


class BaseStream(object):
    """
    A utility base class that we'll use to make sure that our
    streams' respective popNext() methods make use of Python's
    __iter__ magic method
    """
    def __init__(self):
        self.iter = self.__iter__()

    def popNext(self):
        try:
            return self.iter.next()
        except StopIteration:
            return None


class Primes(BaseStream):
    """
    Primes - returns a stream of ordered prime numbers 
    """
    def __iter__(self):
        current = 2
        not_prime = defaultdict(list)
        while True:
            if current in not_prime:
                factors = not_prime.pop(current)
            else:
                yield current
                factors = [current]
            for f in factors:
                not_prime[current + f].append(f)
            current += 1

"""
Higher-order functions to test your classes:
"""
def map(fn, stream):
    """
    map(fn, stream)

    Returns a stream where fn has been applied to each element.
    """
    for s in stream:
        yield fn(s)


def filter(fn, stream):
    """
    filter(fn, stream)

    Returns a stream containing only the elements of the stream for
    which fn returns True.
    """
    for s in stream:
        if fn(s):
            yield s


def prefixReduce(fn, stream, init):
    """
    prefixReduce(fn, stream, init)

    Where fn(x,y) is a function to perform a reduction across the 
    stream, returns a stream where the nth element is the result of
    combining the first n elements of the input stream using fn.
    """
    y = init
    for s in stream:
        x, y = y, s
        y = fn(x, y)
        yield y