Python 中的 yield

yield 優化內存佔用

有這樣一個例子:

def return_test():
   return [i for i in range(10)]

list = return_test()
print(list)

return_test() 返回一個 list,然後打印這個 list。
貌似沒有什麼問題。
那麼再來看下面這個例子:

def return_test():
   return [i for i in range(100000)]

list = return_test()
print(list)

和上面例子唯一不同的是 return_test() 返回的 list 長度變成了 10 萬,數據量變多了。
由於 list 是要存儲在內存中的,遇到海量的數據就會非常的佔用內存。
這時候 yield 就派上用場了。

def yield_test():
    print("yield_test start.")
    for i in range(10):
        print("for loop start", i)
        yield i
        print("for loop end", i)
    print("yield_test end.")

generator = yield_test()
print(type(generator))
for item in generator:
    print("generator", item)

帶有 yield 的函數不再是一個普通函數,而是一個生成器 generator。
generator 和 list 一樣,可以迭代。

其實 yield 可以理解爲和 return 一樣,只是會記住返回的位置,下次迭代的時候會從返回的位置開始繼續往下執行。

generator = yield_test() 執行完之後 yield_test 方法並沒有執行,而是等到 generator 被迭代的時候纔開始執行。

generator 有一個 next() 方法,每次被迭代都是調用 next() 方法,也就是每次都執行一次 for 循環。

由於不是把所有值先準備好放內存中,而是每次迭代的時候生成,所以內存佔用大大的減少了。

yield 實現協程

典型的例子就是生產者消費者

def consumer():
    print("--【consumer】start --")
    response = None
    while True:
        print('【consumer】中斷執行,保存上下文')
        print("--> 切換到 produce")
        # n = yield response 分兩步。1. yield response 2. n =
        # 第一步 yield 把值返回給生產者之後即停止
        # 當調用 c.send(n) 繼續執行時接收 send 的參數值賦給 n
        n = yield response
        print('【consumer】獲取上下文,繼續往下執行')
        if not n:
            return
        print(f"【consumer】: 消費 {n} ..")
        response = "OK"

def produce(c):
    print("--【produce】start --")

    print("【produce】啓動生成器...")
    print("--> 切換到 consumer")
    c.send(None)  # 啓動生成器必須傳入 None
    print("【produce】啓動生成器成功")

    n = 0
    while n < 5:
        n += 1
        print(f"【produce】生產 {n} ..")
        print("--> 切換到 consumer")
        # 把 n 發送給消費者,n = yield response 這段代碼的 n 會接收發送的值
        # r 即爲 yield response 的 response 的值
        r = c.send(n)
        print(f"【produce】消費者 return {r} ..")

    # 關閉生成器
    c.close()

if __name__ == "__main__":
    c = consumer()
    produce(c)

consumer() 獲取生成器。
調用 send(None) 啓動生成器,consumer 方法的代碼開始執行,遇到 yield 掛起 consumer 繼續執行 produce。
調用 send(n) 從上次掛起的地方繼續執行 consumer 方法的代碼。
調用 close() 關閉生成器。
運行結果如下:

--【produce】start --
【produce】啓動生成器...
--> 切換到 consumer
--【consumer】start --
【consumer】中斷執行,保存上下文
--> 切換到 produce
【produce】啓動生成器成功
【produce】生產 1 ..
--> 切換到 consumer
【consumer】獲取上下文,繼續往下執行
【consumer】: 消費 1 ..
【consumer】中斷執行,保存上下文
--> 切換到 produce
【produce】消費者 return OK ..
【produce】生產 2 ..
--> 切換到 consumer
【consumer】獲取上下文,繼續往下執行
【consumer】: 消費 2 ..
【consumer】中斷執行,保存上下文
--> 切換到 produce
【produce】消費者 return OK ..
【produce】生產 3 ..
--> 切換到 consumer
【consumer】獲取上下文,繼續往下執行
【consumer】: 消費 3 ..
【consumer】中斷執行,保存上下文
--> 切換到 produce
【produce】消費者 return OK ..
【produce】生產 4 ..
--> 切換到 consumer
【consumer】獲取上下文,繼續往下執行
【consumer】: 消費 4 ..
【consumer】中斷執行,保存上下文
--> 切換到 produce
【produce】消費者 return OK ..
【produce】生產 5 ..
--> 切換到 consumer
【consumer】獲取上下文,繼續往下執行
【consumer】: 消費 5 ..
【consumer】中斷執行,保存上下文
--> 切換到 produce
【produce】消費者 return OK ..
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章