迭代器实现

要求: 写一个生成素数的生成器, 但生成一定数量之后就会停止

In : p = prime(3)

In : next(p)
Out: 2

In : next(p)
Out: 3

In : next(p)
Out: 5

In : next(p)
---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
<ipython-input-6-aa41e7f2fa96> in <module>()
----> 1 next(p)

StopIteration:

In : list(prime(5))
Out: [2, 3, 5, 7, 11]

代码

# -*- coding: utf-8 -*-
"""
Created on Fri Apr 27 20:22:20 2018

@author: hong
"""

class Prime:
    def __init__(self,num):
        self.prime= 0
        self.num = num
        #内部计数
        self.count = 0

    def __iter__(self):
        return self
    
    def is_prime(self,n):
        if n in (2,3):
            return True
        if n%2 == 0 or self.prime==1:
            return False
        #2,3是素数,1和偶数不是素数上面已经处理,下面判断剩下的奇数
        i=3
        while True:
            if n%i == 0:
                return False
            i+=2
            if i*i>n:
                return True
            
    def __next__(self):
        flag = True
        while flag:
            self.prime+=1
            if self.count>=self.num:
                raise StopIteration();
            if self.is_prime(self.prime):
                flag = False
                self.count+=1
        return self.prime