迭代器與生成器

迭代器與生成器

迭代是 Python 最強大的功能之一。初看起來，你可能會簡單的認爲迭代只不過是處理序列中元素的一種方法。 然而，絕非僅僅就是如此，還有很多你可能不知道的， 比如創建你自己的迭代器對象，在 itertools 模塊中使用有用的迭代模式，構造生成器函數等等。 這一章目的就是向你展示跟迭代有關的各種常見問題。

手動遍歷迭代器

問題

你想遍歷一個可迭代對象中的所有元素，但是卻不想使用 for 循環。

解決方案

爲了手動的遍歷可迭代對象，使用 next()函數並在代碼中捕獲 StopIteration 異常。 比如，下面的例子手動讀取一個文件中的所有行：

def manual_iter():
    with open('/etc/passwd') as f:
        try:
            while True:
                line = next(f)
                print(line, end='')
        except StopIteration:
            pass

通常來講， StopIteration 用來指示迭代的結尾。 然而，如果你手動使用上面演示的 next()函數的話，你還可以通過返回一個指定值來標記結尾，比如 None。 下面是示例：

with open('/etc/passwd') as f:
    while True:
        line = next(f)
        if line is None:
            break
        print(line, end='')

討論

大多數情況下，我們會使用 for 循環語句用來遍歷一個可迭代對象。 但是，偶爾也需要對迭代做更加精確的控制，這時候瞭解底層迭代機制就顯得尤爲重要了。

下面的交互示例向我們演示了迭代期間所發生的基本細節：

>>> items = [1, 2, 3]
>>> # Get the iterator
>>> it = iter(items) # Invokes items.__iter__()
>>> # Run the iterator
>>> next(it) # Invokes it.__next__()
1
>>> next(it)
2
>>> next(it)
3
>>> next(it)
Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
StopIteration
>>>

本章接下來幾小節會更深入的講解迭代相關技術，前提是你先要理解基本的迭代協議機制。 所以確保你已經把這章的內容牢牢記在心中。

代理迭代

問題

你構建了一個自定義容器對象，裏面包含有列表、元組或其他可迭代對象。 你想直接在你的這個新容器對象上執行迭代操作。

解決方案

實際上你只需要定義一個 iter() 方法，將迭代操作代理到容器內部的對象上去。比如：

class Node:
    def __init__(self, value):
        self._value = value
        self._children = []

    def __repr__(self):
        return 'Node({!r})'.format(self._value)

    def add_child(self, node):
        self._children.append(node)

    def __iter__(self):
        return iter(self._children)

# Example
if __name__ == '__main__':
    root = Node(0)
    child1 = Node(1)
    child2 = Node(2)
    root.add_child(child1)
    root.add_child(child2)
    # Outputs Node(1), Node(2)
    for ch in root:
        print(ch)

在上面代碼中， iter()方法只是簡單的將迭代請求傳遞給內部的 _children屬性。

討論

Python 的迭代器協議需要iter()方法返回一個實現了 next()方法的迭代器對象。 如果你只是迭代遍歷其他容器的內容，你無須擔心底層是怎樣實現的。你所要做的只是傳遞迭代請求既可。

這裏的iter() 函數的使用簡化了代碼，iter(s)只是簡單的通過調用 s.iter()方法來返回對應的迭代器對象， 就跟 len(s) 會調用 s.len()原理是一樣的。

使用生成器創建新的迭代模式

問題

你想實現一個自定義迭代模式，跟普通的內置函數比如 range(), reversed()不一樣。

解決方案

如果你想實現一種新的迭代模式，使用一個生成器函數來定義它。 下面是一個生產某個範圍內浮點數的生成器：

def frange(start, stop, increment):
    x = start
    while x < stop:
        yield x
        x += increment

爲了使用這個函數， 你可以用 for 循環迭代它或者使用其他接受一個可迭代對象的函數(比如 sum(), list() 等)。示例如下：

>>> for n in frange(0, 4, 0.5):
...     print(n)
...
0
0.5
1.0
1.5
2.0
2.5
3.0
3.5
>>> list(frange(0, 1, 0.125))
[0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875]
>>>

討論

一個函數中需要有一個 yield 語句即可將其轉換爲一個生成器。 跟普通函數不同的是，生成器只能用於迭代操作。 下面是一個實驗，向你展示這樣的函數底層工作機制：

>>> def countdown(n):
...     print('Starting to count from', n)
...     while n > 0:
...         yield n
...         n -= 1
...     print('Done!')
...

>>> # Create the generator, notice no output appears
>>> c = countdown(3)
>>> c
<generator object countdown at 0x1006a0af0>

>>> # Run to first yield and emit a value
>>> next(c)
Starting to count from 3
3

>>> # Run to the next yield
>>> next(c)
2

>>> # Run to next yield
>>> next(c)
1

>>> # Run to next yield (iteration stops)
>>> next(c)
Done!
Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
StopIteration
>>>

一個生成器函數主要特徵是它只會迴應在迭代中使用到的 next 操作。 一旦生成器函數返回退出，迭代終止。我們在迭代中通常使用的 for 語句會自動處理這些細節，所以你無需擔心。

實現迭代器協議

問題

你想構建一個能支持迭代操作的自定義對象，並希望找到一個能實現迭代協議的簡單方法。

解決方案

目前爲止，在一個對象上實現迭代最簡單的方式是使用一個生成器函數。 在上面小節中，使用 Node 類來表示樹形數據結構。你可能想實現一個以深度優先方式遍歷樹形節點的生成器。 下面是代碼示例：

class Node:
    def __init__(self, value):
        self._value = value
        self._children = []

    def __repr__(self):
        return 'Node({!r})'.format(self._value)

    def add_child(self, node):
        self._children.append(node)

    def __iter__(self):
        return iter(self._children)

    def depth_first(self):
        yield self
        for c in self:
            yield from c.depth_first()

# Example
if __name__ == '__main__':
    root = Node(0)
    child1 = Node(1)
    child2 = Node(2)
    root.add_child(child1)
    root.add_child(child2)
    child1.add_child(Node(3))
    child1.add_child(Node(4))
    child2.add_child(Node(5))

    for ch in root.depth_first():
        print(ch)
    # Outputs Node(0), Node(1), Node(3), Node(4), Node(2), Node(5)

在這段代碼中，depth_first()方法簡單直觀。 它首先返回自己本身並迭代每一個子節點並 通過調用子節點的 depth_first() 方法(使用yield from語句)返回對應元素。

討論

Python 的迭代協議要求一個 iter() 方法返回一個特殊的迭代器對象， 這個迭代器對象實現了 next() 方法並通過 StopIteration 異常標識迭代的完成。 但是，實現這些通常會比較繁瑣。 下面我們演示下這種方式，如何使用一個關聯迭代器類重新實現 depth_first() 方法：

class Node2:
    def __init__(self, value):
        self._value = value
        self._children = []

    def __repr__(self):
        return 'Node({!r})'.format(self._value)

    def add_child(self, node):
        self._children.append(node)

    def __iter__(self):
        return iter(self._children)

    def depth_first(self):
        return DepthFirstIterator(self)

class DepthFirstIterator(object):
    '''
    Depth-first traversal
    '''

    def __init__(self, start_node):
        self._node = start_node
        self._children_iter = None
        self._child_iter = None

    def __iter__(self):
        return self

    def __next__(self):
        # Return myself if just started; create an iterator for children
        if self._children_iter is None:
            self._children_iter = iter(self._node)
            return self._node
        # If processing a child, return its next item
        elif self._child_iter:
            try:
                nextchild = next(self._child_iter)
                return nextchild
            except StopIteration:
                self._child_iter = None
                return next(self)
        # Advance to the next child and start its iteration
        else:
            self._child_iter = next(self._children_iter).depth_first()
            return next(self)

DepthFirstIterator類和上面使用生成器的版本工作原理類似， 但是它寫起來很繁瑣，因爲迭代器必須在迭代處理過程中維護大量的狀態信息。 坦白來講，沒人願意寫這麼晦澀的代碼。將你的迭代器定義爲一個生成器後一切迎刃而解。

反向迭代

問題

你想反方向迭代一個序列

#####解決方案
使用內置的 reversed()函數，比如：

>>> a = [1, 2, 3, 4]
>>> for x in reversed(a):
...     print(x)
...
4
3
2
1

反向迭代僅僅當對象的大小可預先確定或者對象實現了 reversed() 的特殊方法時才能生效。 如果兩者都不符合，那你必須先將對象轉換爲一個列表才行，比如：

# Print a file backwards
f = open('somefile')
for line in reversed(list(f)):
    print(line, end='')

要注意的是如果可迭代對象元素很多的話，將其預先轉換爲一個列表要消耗大量的內存。

討論

很多程序員並不知道可以通過在自定義類上實現reversed() 方法來實現反向迭代。比如：

class Countdown:
    def __init__(self, start):
        self.start = start

    # Forward iterator
    def __iter__(self):
        n = self.start
        while n > 0:
            yield n
            n -= 1

    # Reverse iterator
    def __reversed__(self):
        n = 1
        while n <= self.start:
            yield n
            n += 1

for rr in reversed(Countdown(30)):
    print(rr)
for rr in Countdown(30):
    print(rr)

定義一個反向迭代器可以使得代碼非常的高效， 因爲它不再需要將數據填充到一個列表中然後再去反向迭代這個列表。

帶有外部狀態的生成器函數

問題

你想定義一個生成器函數，但是它會調用某個你想暴露給用戶使用的外部狀態值。

解決方案

如果你想讓你的生成器暴露外部狀態給用戶， 別忘了你可以簡單的將它實現爲一個類，然後把生成器函數放到 iter() 方法中過去。比如：

from collections import deque

class linehistory:
    def __init__(self, lines, histlen=3):
        self.lines = lines
        self.history = deque(maxlen=histlen)

    def __iter__(self):
        for lineno, line in enumerate(self.lines, 1):
            self.history.append((lineno, line))
            yield line

    def clear(self):
        self.history.clear()

爲了使用這個類，你可以將它當做是一個普通的生成器函數。 然而，由於可以創建一個實例對象，於是你可以訪問內部屬性值， 比如 history 屬性或者是 clear() 方法。代碼示例如下：

with open('somefile.txt') as f:
    lines = linehistory(f)
    for line in lines:
        if 'python' in line:
            for lineno, hline in lines.history:
                print('{}:{}'.format(lineno, hline), end='')

討論

關於生成器，很容易掉進函數無所不能的陷阱。 如果生成器函數需要跟你的程序其他部分打交道的話(比如暴露屬性值，允許通過方法調用來控制等等)， 可能會導致你的代碼異常的複雜。 如果是這種情況的話，可以考慮使用上面介紹的定義類的方式。 在 iter()方法中定義你的生成器不會改變你任何的算法邏輯。 由於它是類的一部分，所以允許你定義各種屬性和方法來供用戶使用。

一個需要注意的小地方是，如果你在迭代操作時不使用 for 循環語句，那麼你得先調用 iter()函數。比如：

>>> f = open('somefile.txt')
>>> lines = linehistory(f)
>>> next(lines)
Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
TypeError: 'linehistory' object is not an iterator

>>> # Call iter() first, then start iterating
>>> it = iter(lines)
>>> next(it)
'hello world\n'
>>> next(it)
'this is a test\n'
>>>

迭代器切片

問題

你想得到一個由迭代器生成的切片對象，但是標準切片操作並不能做到。

解決方案

函數 itertools.islice()正好適用於在迭代器和生成器上做切片操作。比如：

>>> def count(n):
...     while True:
...         yield n
...         n += 1
...
>>> c = count(0)
>>> c[10:20]
Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
TypeError: 'generator' object is not subscriptable

>>> # Now using islice()
>>> import itertools
>>> for x in itertools.islice(c, 10, 20):
...     print(x)
...
10
11
12
13
14
15
16
17
18
19
>>>

討論

迭代器和生成器不能使用標準的切片操作，因爲它們的長度事先我們並不知道(並且也沒有實現索引)。 函數 islice()返回一個可以生成指定元素的迭代器，它通過遍歷並丟棄直到切片開始索引位置的所有元素。 然後纔開始一個個的返回元素，並直到切片結束索引位置。

這裏要着重強調的一點是 islice() 會消耗掉傳入的迭代器中的數據。 必須考慮到迭代器是不可逆的這個事實。 所以如果你需要之後再次訪問這個迭代器的話，那你就得先將它裏面的數據放入一個列表中。

跳過可迭代對象的開始部分

問題

你想遍歷一個可迭代對象，但是它開始的某些元素你並不感興趣，想跳過它們。

解決方案

itertools 模塊中有一些函數可以完成這個任務。 首先介紹的是 itertools.dropwhile()函數。使用時，你給它傳遞一個函數對象和一個可迭代對象。 它會返回一個迭代器對象，丟棄原有序列中直到函數返回 True 之前的所有元素，然後返回後面所有元素。

爲了演示，假定你在讀取一個開始部分是幾行註釋的源文件。比如：

>>> with open('/etc/passwd') as f:
... for line in f:
...     print(line, end='')
...
##
# User Database
#
# Note that this file is consulted directly only when the system is running
# in single-user mode. At other times, this information is provided by
# Open Directory.
...
##
nobody:*:-2:-2:Unprivileged User:/var/empty:/usr/bin/false
root:*:0:0:System Administrator:/var/root:/bin/sh
...
>>>

如果你想跳過開始部分的註釋行的話，可以這樣做：

>>> from itertools import dropwhile
>>> with open('/etc/passwd') as f:
...     for line in dropwhile(lambda line: line.startswith('#'), f):
...         print(line, end='')
...
nobody:*:-2:-2:Unprivileged User:/var/empty:/usr/bin/false
root:*:0:0:System Administrator:/var/root:/bin/sh
...
>>>

這個例子是基於根據某個測試函數跳過開始的元素。 如果你已經明確知道了要跳過的元素的個數的話，那麼可以使用 itertools.islice() 來代替。比如：

>>> from itertools import islice
>>> items = ['a', 'b', 'c', 1, 4, 10, 15]
>>> for x in islice(items, 3, None):
...     print(x)
...
1
4
10
15
>>>

在這個例子中， islice() 函數最後那個 None 參數指定了你要獲取從第3個到最後的所有元素， 如果 None 和3的位置對調，意思就是僅僅獲取前三個元素恰恰相反， (這個跟切片的相反操作 [3:] 和 [:3]原理是一樣的)。

討論

函數 dropwhile() 和 islice() 其實就是兩個幫助函數，爲的就是避免寫出下面這種冗餘代碼：

with open('/etc/passwd') as f:
    # Skip over initial comments
    while True:
        line = next(f, '')
        if not line.startswith('#'):
            break

    # Process remaining lines
    while line:
        # Replace with useful processing
        print(line, end='')
        line = next(f, None)

跳過一個可迭代對象的開始部分跟通常的過濾是不同的。 比如，上述代碼的第一個部分可能會這樣重寫：

with open('/etc/passwd') as f:
    lines = (line for line in f if not line.startswith('#'))
    for line in lines:
        print(line, end='')

這樣寫確實可以跳過開始部分的註釋行，但是同樣也會跳過文件中其他所有的註釋行。 換句話講，我們的解決方案是僅僅跳過開始部分滿足測試條件的行，在那以後，所有的元素不再進行測試和過濾了。

最後需要着重強調的一點是，本節的方案適用於所有可迭代對象，包括那些事先不能確定大小的， 比如生成器，文件及其類似的對象。

排列組合的迭代

問題

你想迭代遍歷一個集合中元素的所有可能的排列或組合

#####解決方案
itertools 模塊提供了三個函數來解決這類問題。 其中一個是 itertools.permutations()， 它接受一個集合併產生一個元組序列，每個元組由集合中所有元素的一個可能排列組成。 也就是說通過打亂集合中元素排列順序生成一個元組，比如：

>>> items = ['a', 'b', 'c']
>>> from itertools import permutations
>>> for p in permutations(items):
...     print(p)
...
('a', 'b', 'c')
('a', 'c', 'b')
('b', 'a', 'c')
('b', 'c', 'a')
('c', 'a', 'b')
('c', 'b', 'a')
>>>

如果你想得到指定長度的所有排列，你可以傳遞一個可選的長度參數。就像這樣：

>>> for p in permutations(items, 2):
...     print(p)
...
('a', 'b')
('a', 'c')
('b', 'a')
('b', 'c')
('c', 'a')
('c', 'b')
>>>

使用 itertools.combinations()可得到輸入集合中元素的所有的組合。比如：

>>> from itertools import combinations
>>> for c in combinations(items, 3):
...     print(c)
...
('a', 'b', 'c')

>>> for c in combinations(items, 2):
...     print(c)
...
('a', 'b')
('a', 'c')
('b', 'c')

>>> for c in combinations(items, 1):
...     print(c)
...
('a',)
('b',)
('c',)
>>>

對於 combinations() 來講，元素的順序已經不重要了。 也就是說，組合 ('a', 'b')跟 ('b', 'a')其實是一樣的(最終只會輸出其中一個)。

在計算組合的時候，一旦元素被選取就會從候選中剔除掉(比如如果元素’a’已經被選取了，那麼接下來就不會再考慮它了)。 而函數 itertools.combinations_with_replacement()允許同一個元素被選擇多次，比如：

>>> for c in combinations_with_replacement(items, 3):
...     print(c)
...
('a', 'a', 'a')
('a', 'a', 'b')
('a', 'a', 'c')
('a', 'b', 'b')
('a', 'b', 'c')
('a', 'c', 'c')
('b', 'b', 'b')
('b', 'b', 'c')
('b', 'c', 'c')
('c', 'c', 'c')
>>>

討論

這一小節我們向你展示的僅僅是 itertools 模塊的一部分功能。 儘管你也可以自己手動實現排列組合算法，但是這樣做得要花點腦力。 當我們碰到看上去有些複雜的迭代問題時，最好可以先去看看 itertools 模塊。 如果這個問題很普遍，那麼很有可能會在裏面找到解決方案！

序列上索引值迭代

問題

你想在迭代一個序列的同時跟蹤正在被處理的元素索引。

解決方案

內置的 enumerate() 函數可以很好的解決這個問題：

>>> my_list = ['a', 'b', 'c']
>>> for idx, val in enumerate(my_list):
...     print(idx, val)
...
0 a
1 b
2 c

爲了按傳統行號輸出(行號從1開始)，你可以傳遞一個開始參數：

>>> my_list = ['a', 'b', 'c']
>>> for idx, val in enumerate(my_list, 1):
...     print(idx, val)
...
1 a
2 b
3 c

這種情況在你遍歷文件時想在錯誤消息中使用行號定位時候非常有用：

def parse_data(filename):
    with open(filename, 'rt') as f:
        for lineno, line in enumerate(f, 1):
            fields = line.split()
            try:
                count = int(fields[1])
                ...
            except ValueError as e:
                print('Line {}: Parse error: {}'.format(lineno, e))

enumerate() 對於跟蹤某些值在列表中出現的位置是很有用的。 所以，如果你想將一個文件中出現的單詞映射到它出現的行號上去，可以很容易的利用 enumerate()來完成：

word_summary = defaultdict(list)

with open('myfile.txt', 'r') as f:
    lines = f.readlines()

for idx, line in enumerate(lines):
    # Create a list of words in current line
    words = [w.strip().lower() for w in line.split()]
    for word in words:
        word_summary[word].append(idx)

如果你處理完文件後打印 word_summary，會發現它是一個字典(準確來講是一個 defaultdict )， 對於每個單詞有一個 key ，每個 key 對應的值是一個由這個單詞出現的行號組成的列表。 如果某個單詞在一行中出現過兩次，那麼這個行號也會出現兩次， 同時也可以作爲文本的一個簡單統計。

討論

當你想額外定義一個計數變量的時候，使用 enumerate() 函數會更加簡單。你可能會像下面這樣寫代碼：

lineno = 1
for line in f:
    # Process line
    ...
    lineno += 1

但是如果使用 enumerate() 函數來代替就顯得更加優雅了：

for lineno, line in enumerate(f):
    # Process line
    ...

enumerate()函數返回的是一個 enumerate對象實例， 它是一個迭代器，返回連續的包含一個計數和一個值的元組， 元組中的值通過在傳入序列上調用 next()返回。

還有一點可能並不很重要，但是也值得注意， 有時候當你在一個已經解壓後的元組序列上使用 enumerate() 函數時很容易調入陷阱。 你得像下面正確的方式這樣寫：

data = [ (1, 2), (3, 4), (5, 6), (7, 8) ]

# Correct!
for n, (x, y) in enumerate(data):
    ...
# Error!
for n, x, y in enumer

同時迭代多個序列

問題

你想同時迭代多個序列，每次分別從一個序列中取一個元素。

解決方案

爲了同時迭代多個序列，使用 zip() 函數。比如：

>>> xpts = [1, 5, 4, 2, 10, 7]
>>> ypts = [101, 78, 37, 15, 62, 99]
>>> for x, y in zip(xpts, ypts):
...     print(x,y)
...
1 101
5 78
4 37
2 15
10 62
7 99
>>>

zip(a, b)會生成一個可返回元組 (x, y) 的迭代器，其中 x 來自 a，y 來自 b。 一旦其中某個序列到底結尾，迭代宣告結束。 因此迭代長度跟參數中最短序列長度一致。

>>> a = [1, 2, 3]
>>> b = ['w', 'x', 'y', 'z']
>>> for i in zip(a,b):
...     print(i)
...
(1, 'w')
(2, 'x')
(3, 'y')
>>>

如果這個不是你想要的效果，那麼還可以使用 itertools.zip_longest() 函數來代替。比如：

>>> from itertools import zip_longest
>>> for i in zip_longest(a,b):
...     print(i)
...
(1, 'w')
(2, 'x')
(3, 'y')
(None, 'z')

>>> for i in zip_longest(a, b, fillvalue=0):
...     print(i)
...
(1, 'w')
(2, 'x')
(3, 'y')
(0, 'z')
>>>

討論

當你想成對處理數據的時候 zip() 函數是很有用的。 比如，假設你頭列表和一個值列表，就像下面這樣：

headers = ['name', 'shares', 'price']
values = ['ACME', 100, 490.1]

使用 zip() 可以讓你將它們打包並生成一個字典：

s = dict(zip(headers,values))

或者你也可以像下面這樣產生輸出：

for name, val in zip(headers, values):
    print(name, '=', val)

雖然不常見，但是 zip()可以接受多於兩個的序列的參數。 這時候所生成的結果元組中元素個數跟輸入序列個數一樣。比如;

>>> a = [1, 2, 3]
>>> b = [10, 11, 12]
>>> c = ['x','y','z']
>>> for i in zip(a, b, c):
...     print(i)
...
(1, 10, 'x')
(2, 11, 'y')
(3, 12, 'z')
>>>

最後強調一點就是，zip() 會創建一個迭代器來作爲結果返回。 如果你需要將結對的值存儲在列表中，要使用list() 函數。比如：

>>> zip(a, b)
<zip object at 0x1007001b8>
>>> list(zip(a, b))
[(1, 10), (2, 11), (3, 12)]
>>>

不同集合上元素的迭代

問題

你想在多個對象執行相同的操作，但是這些對象在不同的容器中，你希望代碼在不失可讀性的情況下避免寫重複的循環。

解決方案

itertools.chain() 方法可以用來簡化這個任務。 它接受一個可迭代對象列表作爲輸入，並返回一個迭代器，有效的屏蔽掉在多個容器中迭代細節。 爲了演示清楚，考慮下面這個例子：

>>> from itertools import chain
>>> a = [1, 2, 3, 4]
>>> b = ['x', 'y', 'z']
>>> for x in chain(a, b):
... print(x)
...
1
2
3
4
x
y
z
>>>

使用 chain() 的一個常見場景是當你想對不同的集合中所有元素執行某些操作的時候。比如：

# Various working sets of items
active_items = set()
inactive_items = set()

# Iterate over all items
for item in chain(active_items, inactive_items):
    # Process item

這種解決方案要比像下面這樣使用兩個單獨的循環更加優雅，

for item in active_items:
    # Process item
    ...

for item in inactive_items:
    # Process item
    ...

討論

itertools.chain() 接受一個或多個可迭代對象最爲輸入參數。 然後創建一個迭代器，依次連續的返回每個可迭代對象中的元素。 這種方式要比先將序列合併再迭代要高效的多。比如：

# Inefficent
for x in a + b:
    ...

# Better
for x in chain(a, b):
    ...

第一種方案中，a + b操作會創建一個全新的序列並要求a和b的類型一致。 chian() 不會有這一步，所以如果輸入序列非常大的時候會很省內存。 並且當可迭代對象類型不一樣的時候 chain() 同樣可以很好的工作。

創建數據處理管道

問題

你想以數據管道(類似 Unix 管道)的方式迭代處理數據。 比如，你有個大量的數據需要處理，但是不能將它們一次性放入內存中。

解決方案

生成器函數是一個實現管道機制的好辦法。 爲了演示，假定你要處理一個非常大的日誌文件目錄：

foo/
    access-log-012007.gz
    access-log-022007.gz
    access-log-032007.gz
    ...
    access-log-012008
bar/
    access-log-092007.bz2
    ...
    access-log-022008

假設每個日誌文件包含這樣的數據：

124.115.6.12 - - [10/Jul/2012:00:18:50 -0500] "GET /robots.txt ..." 200 71
210.212.209.67 - - [10/Jul/2012:00:18:51 -0500] "GET /ply/ ..." 200 11875
210.212.209.67 - - [10/Jul/2012:00:18:51 -0500] "GET /favicon.ico ..." 404 369
61.135.216.105 - - [10/Jul/2012:00:20:04 -0500] "GET /blog/atom.xml ..." 304 -
...

爲了處理這些文件，你可以定義一個由多個執行特定任務獨立任務的簡單生成器函數組成的容器。就像這樣：

import os
import fnmatch
import gzip
import bz2
import re

def gen_find(filepat, top):
    '''
    Find all filenames in a directory tree that match a shell wildcard pattern
    '''
    for path, dirlist, filelist in os.walk(top):
        for name in fnmatch.filter(filelist, filepat):
            yield os.path.join(path,name)

def gen_opener(filenames):
    '''
    Open a sequence of filenames one at a time producing a file object.
    The file is closed immediately when proceeding to the next iteration.
    '''
    for filename in filenames:
        if filename.endswith('.gz'):
            f = gzip.open(filename, 'rt')
        elif filename.endswith('.bz2'):
            f = bz2.open(filename, 'rt')
        else:
            f = open(filename, 'rt')
        yield f
        f.close()

def gen_concatenate(iterators):
    '''
    Chain a sequence of iterators together into a single sequence.
    '''
    for it in iterators:
        yield from it

def gen_grep(pattern, lines):
    '''
    Look for a regex pattern in a sequence of lines
    '''
    pat = re.compile(pattern)
    for line in lines:
        if pat.search(line):
            yield line

現在你可以很容易的將這些函數連起來創建一個處理管道。 比如，爲了查找包含單詞 python 的所有日誌行，你可以這樣做：

lognames = gen_find('access-log*', 'www')
files = gen_opener(lognames)
lines = gen_concatenate(files)
pylines = gen_grep('(?i)python', lines)
for line in pylines:
    print(line)

如果將來的時候你想擴展管道，你甚至可以在生成器表達式中包裝數據。 比如，下面這個版本計算出傳輸的字節數並計算其總和。

lognames = gen_find('access-log*', 'www')
files = gen_opener(lognames)
lines = gen_concatenate(files)
pylines = gen_grep('(?i)python', lines)
bytecolumn = (line.rsplit(None,1)[1] for line in pylines)
bytes = (int(x) for x in bytecolumn if x != '-')
print('Total', sum(bytes))

討論

以管道方式處理數據可以用來解決各類其他問題，包括解析，讀取實時數據，定時輪詢等。

爲了理解上述代碼，重點是要明白 yield 語句作爲數據的生產者而 for 循環語句作爲數據的消費者。 當這些生成器被連在一起後，每個 yield 會將一個單獨的數據元素傳遞給迭代處理管道的下一階段。 在例子最後部分，sum() 函數是最終的程序驅動者，每次從生成器管道中提取出一個元素。

這種方式一個非常好的特點是每個生成器函數很小並且都是獨立的。這樣的話就很容易編寫和維護它們了。 很多時候，這些函數如果比較通用的話可以在其他場景重複使用。 並且最終將這些組件組合起來的代碼看上去非常簡單，也很容易理解。

使用這種方式的內存效率也不得不提。上述代碼即便是在一個超大型文件目錄中也能工作的很好。 事實上，由於使用了迭代方式處理，代碼運行過程中只需要很小很小的內存。

在調用 gen_concatenate() 函數的時候你可能會有些不太明白。 這個函數的目的是將輸入序列拼接成一個很長的行序列。itertools.chain()函數同樣有類似的功能，但是它需要將所有可迭代對象最爲參數傳入。 在上面這個例子中，你可能會寫類似這樣的語句lines = itertools.chain(*files) ， 使得gen_opener()生成器能被全部消費掉。 但由於 gen_opener()生成器每次生成一個打開過的文件， 等到下一個迭代步驟時文件就關閉了，因此 china() 在這裏不能這樣使用。 上面的方案可以避免這種情況。

gen_concatenate() 函數中出現過 yield from語句，它將 yield操作代理到父生成器上去。 語句 yield from it 簡單的返回生成器 it所產生的所有值。 關於這個我們在4.14小節會有更進一步的描述。

最後還有一點需要注意的是，管道方式並不是萬能的。 有時候你想立即處理所有數據。 然而，即便是這種情況，使用生成器管道也可以將這類問題從邏輯上變爲工作流的處理方式。

David Beazley 在他的 Generator Tricks for Systems Programmers 教程中對於這種技術有非常深入的講解。可以參考這個教程獲取更多的信息。

展開嵌套的序列

問題

你想將一個多層嵌套的序列展開成一個單層列表

#####解決方案
可以寫一個包含 yield from 語句的遞歸生成器來輕鬆解決這個問題。比如：

from collections import Iterable

def flatten(items, ignore_types=(str, bytes)):
    for x in items:
        if isinstance(x, Iterable) and not isinstance(x, ignore_types):
            yield from flatten(x)
        else:
            yield x

items = [1, 2, [3, 4, [5, 6], 7], 8]
# Produces 1 2 3 4 5 6 7 8
for x in flatten(items):
    print(x)

在上面代碼中， isinstance(x, Iterable) 檢查某個元素是否是可迭代的。 如果是的話， yield from 就會返回所有子例程的值。最終返回結果就是一個沒有嵌套的簡單序列了。

額外的參數 ignore_types 和檢測語句 isinstance(x, ignore_types) 用來將字符串和字節排除在可迭代對象外，防止將它們再展開成單個的字符。 這樣的話字符串數組就能最終返回我們所期望的結果了。比如：

>>> items = ['Dave', 'Paula', ['Thomas', 'Lewis']]
>>> for x in flatten(items):
...     print(x)
...
Dave
Paula
Thomas
Lewis
>>>

討論

語句 yield from 在你想在生成器中調用其他生成器作爲子例程的時候非常有用。 如果你不使用它的話，那麼就必須寫額外的for 循環了。比如：

def flatten(items, ignore_types=(str, bytes)):
    for x in items:
        if isinstance(x, Iterable) and not isinstance(x, ignore_types):
            for i in flatten(x):
                yield i
        else:
            yield x

儘管只改了一點點，但是 yield from 語句看上去感覺更好，並且也使得代碼更簡潔清爽。

之前提到的對於字符串和字節的額外檢查是爲了防止將它們再展開成單個字符。 如果還有其他你不想展開的類型，修改參數 ignore_types 即可。

最後要注意的一點是，yield from 在涉及到基於協程和生成器的併發編程中扮演着更加重要的角色。

順序迭代合併後的排序迭代對象

問題

你有一系列排序序列，想將它們合併後得到一個排序序列並在上面迭代遍歷。

解決方案

heapq.merge() 函數可以幫你解決這個問題。比如：

>>> import heapq
>>> a = [1, 4, 7, 10]
>>> b = [2, 5, 6, 11]
>>> for c in heapq.merge(a, b):
...     print(c)
...
1
2
4
5
6
7
10
11

討論

heapq.merge 可迭代特性意味着它不會立馬讀取所有序列。 這就意味着你可以在非常長的序列中使用它，而不會有太大的開銷。 比如，下面是一個例子來演示如何合併兩個排序文件：

with open('sorted_file_1', 'rt') as file1, \
    open('sorted_file_2', 'rt') as file2, \
    open('merged_file', 'wt') as outf:

    for line in heapq.merge(file1, file2):
        outf.write(line)

有一點要強調的是 heapq.merge()需要所有輸入序列必須是排過序的。 特別的，它並不會預先讀取所有數據到堆棧中或者預先排序，也不會對輸入做任何的排序檢測。 它僅僅是檢查所有序列的開始部分並返回最小的那個，這個過程一直會持續直到所有輸入序列中的元素都被遍歷完。

迭代器代替 while 無限循環

問題

你在代碼中使用while循環來迭代處理數據，因爲它需要調用某個函數或者和一般迭代模式不同的測試條件。 能不能用迭代器來重寫這個循環呢？

解決方案

一個常見的 IO 操作程序可能會想下面這樣：

CHUNKSIZE = 8192

def reader(s):
    while True:
        data = s.recv(CHUNKSIZE)
        if data == b'':
            break
        process_data(data)

這種代碼通常可以使用 iter()來代替，如下所示：

def reader2(s):
    for chunk in iter(lambda: s.recv(CHUNKSIZE), b''):
        pass
        # process_data(data)

如果你懷疑它到底能不能正常工作，可以試驗下一個簡單的例子。比如：

>>> import sys
>>> f = open('/etc/passwd')
>>> for chunk in iter(lambda: f.read(10), ''):
...     n = sys.stdout.write(chunk)
...
nobody:*:-2:-2:Unprivileged User:/var/empty:/usr/bin/false
root:*:0:0:System Administrator:/var/root:/bin/sh
daemon:*:1:1:System Services:/var/root:/usr/bin/false
_uucp:*:4:4:Unix to Unix Copy Protocol:/var/spool/uucp:/usr/sbin/uucico
...
>>>

討論

iter 函數一個鮮爲人知的特性是它接受一個可選的 callable 對象和一個標記(結尾)值作爲輸入參數。 當以這種方式使用的時候，它會創建一個迭代器， 這個迭代器會不斷調用 callable對象直到返回值和標記值相等爲止。

這種特殊的方法對於一些特定的會被重複調用的函數很有效果，比如涉及到 I/O 調用的函數。 舉例來講，如果你想從套接字或文件中以數據塊的方式讀取數據，通常你得要不斷重複的執行 read() 或recv() ， 並在後面緊跟一個文件結尾測試來決定是否終止。這節中的方案使用一個簡單的iter()調用就可以將兩者結合起來了。 其中 lambda函數參數是爲了創建一個無參的 callable對象，併爲 recv 或 read()方法提供了 size參數。