[Python進階]Python的iterator和generator

發表於 | 分類於 | |  閱讀次數:
字數統計: 981

剛開始寫Python的時候，你可能對於Python的語法for i in [2, 4, 6, 8, 10]感到驚艷，這種簡潔的語法相比C++或Java來說直觀了許多。然而，你有沒有想過Python在處理for in的時候，背後到底發生了什麼？什麼樣的object可以被放進for in loop呢？

Container和Iterator

Container的概念非常簡單，Python中的一切皆為object，而object的集合就是Container，例如list, set, tuple等等。所有的Container都是iterable（可迭代的）。可迭代是什麼意思呢？你可以想像一下你去水果攤買蘋果，老闆不告訴你庫存情況，每次你只需要跟老闆說「我要一個蘋果」，直到老闆告訴你「蘋果沒了」。Container透過iter()返回一個Iterator，然後我們可以透過next()就像跟老闆要蘋果那樣，一個一個要。

1
2
3
4
5
6
7
8
9
10
11
12
>>> arr = [1, 2, 3]
>>> iter_arr = iter(arr)
>>> next(iter_arr)
1
>>> next(iter_arr)
2
>>> next(iter_arr)
3
>>> next(iter_arr)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

當老闆說沒有蘋果的時候，就會throw exception出來。

Generator

Generator可以理解為懶人版的Container。生成一個Container很簡單，[i for i in range(100000000)]就可以生成一億個int的array，每個元素都會保存在記憶體當中。當然，轉成Iterator後也是，只是他們取用元素的方式不同而已。Generator的使用方法和Iterator類似，都是需要後再取，只是Iterator會預先把所有元素放進記憶體，而Generator會等到需要拿的時候才會把該元素載入記憶體。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import os
import psutil

def show_memory_info(hint):
    pid = os.getpid()
    p = psutil.Process(pid)
    info = p.memory_full_info()
    memory = info.uss / 1024. / 1024
    print('{} memory used: {} MB'.format(hint, memory))

def test_iterator():
    show_memory_info('initing iterator')
    list_1 = [i for i in range(100000000)]
    show_memory_info('after iterator initiated')
    print(sum(list_1))
    show_memory_info('after sum called')

def test_generator():
    show_memory_info('initing generator')
    list_2 = (i for i in range(100000000))
    show_memory_info('after generator initiated')
    print(sum(list_2))
    show_memory_info('after sum called')

test_iterator()
test_generator()
# Output
# initing iterator memory used: 7.28125 MB
# after iterator initiated memory used: 2015.828125 MB
# 4999999950000000
# after sum called memory used: 3869.05078125 MB
# initing generator memory used: 9.5703125 MB
# after generator initiated memory used: 9.58203125 MB
# 4999999950000000
# after sum called memory used: 9.58203125 MB

我們可以看到記憶體驚人的差異，在Iterator和Generator初始記憶體都差不多的情況下，Iterator需要2GB的記憶體，而Generator只需要9.5MB。如果遇到不需要同時在記憶體保存這麼多東西的場景，例如元素總和，可以使用Generator。由上面的範例可以看到，Generator的初始化寫法是(i for i in range(100000000))。

那麼Generator還能怎麼玩呢？例如我們想要驗證一個數學公式{(1+2+3+…+n)^2 = 1^3 + 2^3 + 3^3 + … + n^3}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import time

def generator(k):
   i = 1
   while True:
      yield i ** k
      i+=1

gen_1 = generator(1)
gen_3 = generator(3)
sum_1 = 0
sum_3 = 0
while True:
   sum_1 += next(gen_1)
   sum_2 = sum_1 ** 2
   sum_3 += next(gen_3)
   print("sum_2: %s, sum_3: %s" % (sum_2, sum_3))
   time.sleep(0.5)

yield是Generator獨有的，你可以理解為在next之前，它會被卡在這裡，呼叫next之後yield就會return值出來。你看，有了Generator，我就能一直無限的驗證下去，不用擔心記憶體爆炸。Iterator是一個有限集合，Generator是一個無限集合！

除此之外，Generator也能讓程式碼更加簡潔有力！讓我們看下面的例子，輸入一個array和一個數字，找出該數字在array的位置。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# Iterator
def find_iter(arr, val):
   res = []
   for idx, v in enumerate(arr):
      if v == val:
         res.append(idx)
   return res
   
# Generator
def find_gen(arr, val):
   for idx, v in enumerate(arr):
      if v == val:
         yield idx

arr = [1, 5, 2, 9, 1, 7, 2, 1, 2]
val = 2
print(find_iter(arr, val))
print(list(find_gen(arr, val)))
# Output
# [2, 6, 8]
# [2, 6, 8]

顯然的，Generator清爽多了。

總結

本篇講了Container, Iterator和Generator

  • Container是Iterable的，代表將Container放進for in裡我們可以一個個迭代
  • Generator是一個特殊的Iterator，使用Generator可以寫出更清新，更省資源的程式碼