Python常用内建模块-itertools-猿码集

1. itertools模块介绍

itertools是Python中的一个内置模块，在Python中提供了很多方便快捷的迭代器函数，能够高效地进行迭代操作，对于需要将复杂的迭代操作简单化的开发者非常有用。该模块包括三类函数：

返回可迭代对象的函数

返回迭代器的函数

返回迭代器生成器的函数

1.1 返回可迭代对象的函数

itertools模块中有许多返回可迭代对象的函数，下面我们分别介绍几个常用的函数：

1.1.1 itertools.count(start=0, step=1)

该方法可以创建一个无限的计数器迭代器，从start开始，以step为步长进行计数。


import itertools
counter = itertools.count()
data = [next(counter) for _ in range(5)]
print(data)

输出结果：[0, 1, 2, 3, 4]

当然，我们也可以通过设置start参数进行指定起始值。


import itertools
counter = itertools.count(start=5)
data = [next(counter) for _ in range(5)]
print(data)

输出结果：[5, 6, 7, 8, 9]

我们还可以通过设置step参数进行指定步长。


import itertools
counter = itertools.count(start=0, step=2)
data = [next(counter) for _ in range(5)]
print(data)

输出结果：[0, 2, 4, 6, 8]

1.1.2 itertools.cycle(iterable)

该函数可以对原有数据进行循环迭代，当原数据迭代完成后，会不断重复迭代。


import itertools
counter = itertools.cycle(['A', 'B', 'C'])
data = [next(counter) for _ in range(5)]
print(data)

输出结果：['A', 'B', 'C', 'A', 'B']

1.1.3 itertools.repeat(object, times=None)

该函数可以将一个元素重复迭代指定的次数，如果不指定次数，则会无限地重复迭代。


import itertools
data = [i for i in itertools.repeat(10, 5)]
print(data)

输出结果：[10, 10, 10, 10, 10]

1.2 返回迭代器的函数

接着我们来介绍几个返回迭代器的函数：

1.2.1 itertools.compress(data, selectors)

该函数可以根据selectors的真假值对data进行过滤。


import itertools
data = ['A', 'B', 'C', 'D']
selectors = [True, False, True, False]
filtered_data = itertools.compress(data, selectors)
print(list(filtered_data))

输出结果：['A', 'C']

1.2.2 itertools.dropwhile(predicate, iterable)

该函数会跳过原数据中所有符合条件的元素，直到遇到第一个不符合条件的元素，之后返回剩余部分的迭代器。


import itertools
data = [1, 3, 5, 2, 4, 6]
filtered_data = itertools.dropwhile(lambda x: x<5, data)
print(list(filtered_data))

输出结果：[5, 2, 4, 6]

1.2.3 itertools.takewhile(predicate, iterable)

该函数只返回符合条件的元素，直到遇到第一个不符合条件的元素时就立即停止迭代。


import itertools
data = [1, 3, 5, 2, 4, 6]
filtered_data = itertools.takewhile(lambda x: x<5, data)
print(list(filtered_data))

输出结果：[1, 3]

1.3 返回迭代器生成器的函数

最后我们来介绍几个返回迭代器生成器的函数：

1.3.1 itertools.chain(*iterables)

该函数可以将多个可迭代对象拼接到一起，返回一个生成器，生成器会依次迭代每一个可迭代对象中的元素。


import itertools
data = itertools.chain([1, 2, 3], [4, 5, 6])
print(list(data))

输出结果：[1, 2, 3, 4, 5, 6]

1.3.2 itertools.dropwhile(predicate, iterable)

该函数可以对一个可迭代对象进行分组，分组方式由key函数指定。


import itertools
data = [{'name': 'Tom', 'age':20},
        {'name': 'Jack', 'age':18},
        {'name': 'Lucy', 'age':22},
        {'name': 'Lily', 'age':18}]
grouped_data = itertools.groupby(data, key=lambda x:x['age'])
for key, group in grouped_data:
    print(key, list(group))

输出结果：

20 [{'name': 'Tom', 'age': 20}]
18 [{'name': 'Jack', 'age': 18}, {'name': 'Lily', 'age': 18}]
22 [{'name': 'Lucy', 'age': 22}]

2. itertools源码解析

在介绍完itertools的常用函数后，我们来深入学习一下itertools的源码实现。itertools的源码在Python的官方库中是可以查看的，也可以在网上找到。

2.1 itertools.count源码解析


def count(start=0, step=1):
    # count(10) --> 10 11 12 13 14 ...
    # count(2.5, 0.5) -> 2.5 3.0 3.5 ...
    it = _countiter(start, step)
    return it
class _countiter:
    def __init__(self, value, step):
        self.value = value
        self.step = step
    def __iter__(self):
        return self
    def __next__(self):
        result = self.value
        self.value += self.step
        return result

从源码中可以看出，count函数会返回一个_countiter的实例，而_countiter是一个迭代器，其__next__方法会不断返回value，并将value加上step。

2.2 itertools.cycle源码解析


def cycle(iterable):
    # cycle('ABCD') --> A B C D A B C D A B C D ...
    saved = []
    for element in iterable:
        yield element
        saved.append(element)
    while saved:
        for element in saved:
            yield element

从源码中可以看出，cycle函数返回一个生成器，其yield关键字前后轮流返回输入的可迭代对象的元素，当迭代器迭代完所有元素后，会无限重复地返回原队列的所有元素。

2.3 itertools.repeat源码解析


def repeat(object, times=None):
    # repeat(10, 3) --> 10 10 10
    if times is None:
        while True:
            yield object
    else:
        for i in range(times):
            yield object

从源码中可以看出，repeat函数返回一个生成器，当times为None时，它会不断重复yield object，当times不为None时，它会重复yield object次数为times。

2.4 itertools.chain源码解析


def chain(*iterables):
    # chain('ABC', 'DEF') --> A B C D E F
    for it in iterables:
        for element in it:
            yield element

从源码中可以看出，chain函数返回一个生成器，它会首先迭代第一个可迭代对象中的所有元素，然后迭代第二个，第三个，直到最后一个可迭代对象中的所有元素。

2.5 itertools.groupby源码解析


def groupby(iterable, key=None):
    # groupby('AAAABBBCCDAABBB') -> A B C D A B
    # groupby('ABBCcAD', lambda x: x.upper()) -> A B C A D
    if key is None:
        key = lambda x: x
    it = iter(iterable)
    try:
        first = next(it)
    except StopIteration:
        return
    groupkey = key(first)
    group = [first]
    for item in it:
        itemkey = key(item)
        if itemkey != groupkey:
            yield (groupkey, group)
            groupkey = itemkey
            group = [item]
        else:
            group.append(item)
    yield (groupkey, group)

从源码中可以看出，groupby函数返回一个生成器，它会先迭代一次可迭代对象，并将第一个元素的key值保存起来，然后循环迭代所有元素，当元素的key值与前一个元素的key值不同时，就会生成前一个分组的(groupkey, group)二元组，并将groupkey设置为当前元素的key值，将group保存为新的列表。

3. 总结

itertools模块是Python常用的标准库之一，它提供了非常强大而又方便快捷的迭代函数。通过本文的介绍，我们了解了itertools的一些常用函数，并深入源码解析了其中几个函数的实现原理。学习完本文后，读者可以更加熟练地使用itertools模块对数据进行处理，也可以更深入地理解Python中的迭代机制。

Python常用内建模块-itertools