Functional approach to Sieve of Eratosthenes in Python.

.gitignore

__pycache__/
*.pyc
*.out

Makefile

test: sieve_eratosthenes_haskell.py out.control
	rm *.out
	@for py in jython pypy python python3 \
     python2.4 python2.5 python2.6 python2.7 python3.1 python3.2 ; \
     do	echo $$py ; \
      $$py -mdoctest $< 2>/dev/null && \
      time -p $$py $< 22 >$$py.out && \
      diff -s $$py.out out.control ; done

out.control

2
3
5
7
11
13
17
19
23
29
31
37
41
43
47
53
59
61
67
71
73
79

sieve_eratosthenes_haskell.py

#!/usr/bin/env python
"""Functional approach to Sieve of Eratosthenes in Python.

Translation from Haskell:
    http://en.literateprograms.org/Sieve_of_Eratosthenes_%28Haskell%29
"""
import sys

try:
    from itertools import imap as map
except ImportError: # python 3.x
    map = map

try:
    from itertools import izip as zip
except ImportError: # python 3.x
    zip = zip

try:
    range = xrange # python 2.x
except NameError:
    range = range

try:
    cmp = cmp
except NameError: # python 3.x
    def cmp(a, b): 
        return (a > b) - (a < b)

try:
    head = next
except NameError: # python < 2.6
    def head(iterable):
        for it in iterable:
            return it


def primes():
    """Generate prime numbers indefinitely.

    >>> take(20, primes())
    [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71]
    """
    yield 2
    yield 3
    yield 5
    for p in diff(count(7, 2), nonprimes()):
        yield p

def count(start=0, step=1):
    """`itertools.count()`

    >>> take(5, count(7, 2))
    [7, 9, 11, 13, 15]
    >>> take(5, count(2, 5))
    [2, 7, 12, 17, 22]
    """
    yield start
    for i in count(start+step, step):
        yield i

def nonprimes():
    """Return odd non-primes starting with 9.
    
    >>> take(8, nonprimes())
    [9, 15, 21, 25, 27, 33, 35, 39]
    """
    tail_primes = primes()
    head(tail_primes) # discard 2

    for n in foldr1(merge1, map(multiples, tail_primes)):
        yield n

def merge1(xs, ys):
    """Similar to `merge()` but always starts with the first element of `xs`.

    >>> take(7, merge1(count(1, 2), count(0, 4)))
    [1, 0, 3, 4, 5, 7, 8]
    """
    yield head(xs)
    for i in merge(xs, ys):
        yield i

def foldr1(f, iterable_of_iterables):
    """Lazy `reduce()`.

    :f: (iterable1, iterable2) -> iterable
    :iterable_of_iterables: possibly infinite iterable of iterables

    The result is similar to:
    
        f(... f(f(f(a[0], a[1]), a[2]), a[3]), ...)

    >>> take(8, multiples(3))
    [9, 15, 21, 27, 33, 39, 45, 51]
    >>> take(8, multiples(4))
    [16, 24, 32, 40, 48, 56, 64, 72]
    >>> take(11, foldr1(merge1, iter([multiples(3),
    ...                        multiples(4),
    ...                        iter([17, 31, 31.5, 34])])))
    [9, 15, 16, 17, 21, 24, 27, 31, 31.5, 32, 33]
    """
    def iter_(start_it):
        it = f(start_it, head(iterable_of_iterables))
        yield head(it)
        for item in iter_(it):
            yield item
        for item in it: # for finite iterable_of_iterables
            yield item
    return iter_(head(iterable_of_iterables))

def multiples(p):
    """Return odd multiples of the prime `p` starting with ``p*p``.

    >>> take(5, multiples(3))
    [9, 15, 21, 27, 33]
    >>> take(5, multiples(4))
    [16, 24, 32, 40, 48]
    """
    for n in count(p, 2):
        yield n * p

def merge(xs, ys):
    """Merge two infinite generators.

    1. This function cannot cope with finite lists as it has no
       terminating condition for an empty list.
       
    2. The elements of each list must be sorted when the function is
       called.

    3. This system will eliminate duplicates across the lists, but
       will not be able to cope with lists which have duplicates
       internally.

    >>> take(7, merge(count(1, 2), count(6)))
    [1, 3, 5, 6, 7, 8, 9]
    """
    x, y = head(xs), head(ys)
    if x < y:
        yield x
        for i in merge(xs, cons(y, ys)):
            yield i
    elif x == y:
        yield x
        for i in merge(xs, ys):
            yield i
    elif x > y:
        yield y
        for i in merge(cons(x, xs), ys):
            yield i
    else:
        assert 0

def diff(xs, ys):
    """Remove from `xs` generator all values from `ys` generator.

    1. This function will only work on infinite lists.
    
    2. The supplied lists must be ordered.
    
    3. The interaction between the two lists is not what would be
       generally desired from a more general-purpose function.

    >>> take(7, diff(count(), count(1, 2)))
    [0, 2, 4, 6, 8, 10, 12]
    """
    x, y = head(xs), head(ys)
    if x < y:
        yield x
        for i in diff(xs, cons(y, ys)):
            yield i
    elif x == y:
        for i in diff(xs, ys):
            yield i
    elif x > y:
        for i in diff(cons(x, xs), ys):
            yield i
    else:
        assert 0

def cons(head_, tail):
    """Combine `head_` value & the `tail` generator into a single generator.

    `tail` is either an iterable or `lambda: iterable`

    >>> list(cons(1, [2, 3]))
    [1, 2, 3]
    """
    yield head_
    
    for item in tail:
        yield item

def take(n, iterable):
    """Return the first `n` items from the iterable."""
    return [item for _, item in zip(range(n), iterable)]


def main(argv):
    """
    %prog [number_of_primes]

    Default is 10.
    if `number_of_primes` is 0 then print indefinitely
    """
    if len(sys.argv) == 2:
        n = int(sys.argv[1])
    else:
        n = 10

    for i, p in enumerate(primes()):
        if n == 0 or i < n:
            print(p)
        else:
            break
        

if __name__=="__main__":
    sys.exit(main(sys.argv))