[python] generator-based implementation of Sieve of Eratosthenes (Euler's mod)

common.py

def integers():
    i = 1
    while True:
        yield i
        i += 1

def head(iter, n):
    return [iter.next() for i in xrange(n)]

sieve1.py

from itertools import tee as fission, chain, imap
from common import *

def scale(iter, n):
    return imap(lambda x: n*x, iter)

def remove(iter, iter_unwanted):
    """An extremely naive way of remove elements in `iter_unwanted` from
    `iter`."""

    curr = iter_unwanted.next()
    for i in iter:
        if i == curr:
            curr = iter_unwanted.next()
            continue
        yield i

def sieve():
    iter = integers()

    # skip 1
    iter.next()

    while True:
        p = iter.next()
        yield p

        # needed here otherwise remove() and scale() advance the same generator
        iter, iter2 = fission(iter)
        # Euler's modification - remove p multiplied by the numbers we have now
        iter = remove(iter2, scale(chain([p], iter), p))

sieve2.py

from itertools import ifilter
from common import *


class prime_filter(object):
    def __init__(self, p):
        self.p = p
        self.curr = p

    def __call__(self, x):
        while x > self.curr:
            self.curr += self.p
        is_unwanted = x == self.curr
        if is_unwanted:
            self.curr += self.p
        return not is_unwanted


def sieve():
    iter = integers()

    # skip 1
    iter.next()

    while True:
        p = iter.next()
        yield p

        iter = ifilter(prime_filter(p), iter)