Python 实现 head()

zltzlt

Python 实现 head()

语法

head(iterable, n)

描述

以生成器的形式返回 iterable 的前 n 个元素。如果 iterable 中的元素个数 < n，则返回 iterable 中的所有元素。

参数

参数	描述
iterable	可迭代对象。
n	必须是整数。

返回值

生成器。

例子

>>> def gen():
        yield 1
        yield 2
        yield 3
        yield 4
        yield 5

        
>>> g1 = gen()
>>> list(head(g1, 3))
[1, 2, 3]
>>> g2 = gen()
>>> list(head(g2, 1))
[1]
>>> g3 = gen()
>>> list(head(g3, 8))
[1, 2, 3, 4, 5]

代码

def head(iterable, n):
    """
    head(iterable, n) -> [generator object]
    以生成器的形式返回 iterable 的前 n 个元素
    """
    it = iter(iterable)
    for i in range(n):
        try:
            yield next(it)
        except StopIteration:
            break

永恒的蓝色梦想

class head:
    def __init__(self,iterable,n,/):
        if isinstance(n,int) and n>=0:
            self.__n=n
        else:
            raise TypeError(f"'{n.__class__.__name__}' object cannot be interpreted as an integer")

        self.__it=iter(iterable)

    def __iter__(self):
        return self

    def __next__(self):
        if self.__n:
            self.__n-=1
            return self.__it.__next__()

        raise StopIteration