生成器和迭代器

"""
生成器
- 生成器表达式：(x ** 2 for x in range(1, 11))
- yield关键字
迭代器
- __iter__：返回迭代器对象
- __next__：取下一个迭代值
"""
from math import sqrt


def is_prime(num):
    """判断是不是素数"""
    for factor in range(2, int(sqrt(num)) + 1):
        if num % factor == 0:
            return False
    return True


def prime(num):
    """素数生成器"""
    for val in range(2, num + 1):
        if is_prime(val):
            yield val


class Prime(object):
    """素数迭代器"""
    pass


class Fib(object):
    """斐波拉切数迭代器"""

    def __init__(self, num):
        self.num = num
        self.a, self.b = 0, 1
        self.idx = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.idx < self.num:
            self.a, self.b = self.b, self.a + self.b
            self.idx += 1
            return self.a
        raise StopIteration()


def fib(num):
    """斐波拉切数生成器"""
    a, b = 0, 1
    for _ in range(num):
        a, b = b, a + b
        yield a


def main():
    """主函数"""
    print('fibonacci number'.center(80, '-'))
    for val in Fib(20):
        print(val, end=' ')
    print('\n', '-' * 80, sep='')
    print('prime number'.center(80, '-'))
    for val in prime(100):
        print(val, end=' ')
    print('\n', '-' * 80, sep='')


if __name__ == '__main__':
    main()