generatorで素数

プログラミングHaskellの

take 100 primes

みたいな書き方にいたく感動したので。
とりあえずitertoolの中にtakewhileというのを見つけたのでそれに合う書き方をしてみた。

takewhile(lambda x: x < 100, sieve())

このように書いて100以下の素数リストを手に入れたい。

from itertools import takewhile,ifilter

def sieve():
    x = [2]
    i=2
    while True:
        i += 1
        if len(list(ifilter(lambda k: i%k==0, x))) == 0:
            x.append(i)
            yield i

l = takewhile(lambda x: x < 100, sieve())
print list(l)

所与のtake 100 primesを書こうと5分ぐらい考えたのだけど頭がこんがらがってきたのでここまで。generatorもあんまり理解しないで使ってるだけだしちょっとよくない。
あと検索してみたらgeneratorで素数な人はすでにたくさんいた。
ifilterに加えてcountも使ったほうがよさそう。

参考

Python - generatorで素数生成