Mercurial

#!/usr/bin/env python

"""Fun with generators. Corresponding Haskell implementation:

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

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)))
author	test
date	Thu, 01 Jan 1970 00:00:00 +0000
parents
children
rev	line source
test@0	1 #!/usr/bin/env python
test@0	2
test@0	3 """Fun with generators. Corresponding Haskell implementation:
test@0	4
test@0	5 primes = 2 : sieve [3, 5..]
test@0	6 where sieve (p:ns) = p : sieve [n \| n <- ns, mod n p /= 0]
test@0	7 """
test@0	8
test@0	9 from itertools import dropwhile, ifilter, islice, count, chain
test@0	10
test@0	11 def primes():
test@0	12 """Generate all primes."""
test@0	13 def sieve(ns):
test@0	14 p = ns.next()
test@0	15 # It is important to yield here in order to stop the
test@0	16 # infinite recursion.
test@0	17 yield p
test@0	18 ns = ifilter(lambda n: n % p != 0, ns)
test@0	19 for n in sieve(ns):
test@0	20 yield n
test@0	21
test@0	22 odds = ifilter(lambda i: i % 2 == 1, count())
test@0	23 return chain([2], sieve(dropwhile(lambda n: n < 3, odds)))
test@0	24
test@0	25 if __name__ == "__main__":
test@0	26 import sys
test@0	27 try:
test@0	28 n = int(sys.argv[1])
test@0	29 except (ValueError, IndexError):
test@0	30 n = 10
test@0	31 p = primes()
test@0	32 print "The first %d primes: %s" % (n, list(islice(p, n)))
test@0	33