Mercurial

#!/usr/bin/env python
"""Fun with generators. Corresponding Haskell implementation:

"""Fun with generators. Corresponding Haskell implementation:

primes = 2 : sieve [3, 5..]
    where sieve (p:ns) = p : sieve [n | n <- ns, mod n p /= 0]
"""
primes = 2 : sieve [3, 5..]
    where sieve (p:ns) = p : sieve [n | n <- ns, mod n p /= 0]
"""

from itertools import dropwhile, ifilter, islice, count, chain

from itertools import dropwhile, ifilter, islice, count, chain

def primes():
    """Generate all primes."""
    def sieve(ns):
        p = ns.next()
        # It is important to yield *here* in order to stop the
        # infinite recursion.
        yield p
        ns = ifilter(lambda n: n % p != 0, ns)
        for n in sieve(ns):
            yield n

    odds = ifilter(lambda i: i % 2 == 1, count())
    return chain([2], sieve(dropwhile(lambda n: n < 3, odds)))

if __name__ == "__main__":
    import sys
    try:
        n = int(sys.argv[1])
    except (ValueError, IndexError):
        n = 10
    p = primes()
    print "The first %d primes: %s" % (n, list(islice(p, n)))