熟悉 Golang 的同學(xué)都知道，Golang 里面有一個關(guān)鍵詞叫做defer，它可以實(shí)現(xiàn)延遲調(diào)用。

實(shí)際上在 Python 里面也有相關(guān)的語法，那就是contextlib.ExitStack。

我們來看這樣一個場景：

我有一個函數(shù)parse，它的作用是從 Redis 中持續(xù)讀入數(shù)據(jù)，并寫入到MongoDB 中。示例代碼如下：

1. import json 
2. import redis 
3. import pymongo 
4.  
5. client = redis.Redis() 
6. handler = pymongo.MongoClient().test.data 
7.  
8. def parse(): 
9.     data = client.lpop('test') 
10.     if not data: 
11.         return 
12.     handler.insert_one(json.loads(data)) 

但現(xiàn)在我想增加一個需求，當(dāng)Redis 讀取結(jié)束或者讀取數(shù)據(jù)報錯的時候，能把當(dāng)前的時間也寫入到MongoDB 中。

那么代碼可能變成下面這樣：

1. import json 
2. import redis 
3. import datetime 
4. import pymongo 
5.  
6. client = redis.Redis() 
7. handler = pymongo.MongoClient().test.data 
8.  
9. def parse(): 
10.     while True: 
11.         try: 
12.             data = client.lpop('test') 
13.             if not data: 
14.                 handler.insert_one({'finished': True, 'ts': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}) 
15.                 return 
16.             handler.insert_one(json.loads(data)) 
17.         except Exception: 
18.             handler.insert_one({'finished': True, 'ts': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'}) 

可以看到，代碼變得很難看了。

現(xiàn)在，我們可以使用延遲調(diào)用來讓代碼變得更好看。

要實(shí)現(xiàn)這個目的，就可以開始使用ExitStack了。它可以注冊多個回調(diào)函數(shù)，在退出上下文縮進(jìn)時執(zhí)行。

我們先來看一個簡單的例子：

1. import contextlib 
2.  
3. def callback_1(): 
4.     print('我是第一個回調(diào)函數(shù)') 
5.  
6. def callback_2(x): 
7.     print(f'我是第二個回調(diào)函數(shù)，傳入?yún)?shù)：{x}') 
8.  
9.  
10. with contextlib.ExitStack() as stack: 
11.     stack.callback(callback_1) 
12.     stack.callback(callback_2, 100) 
13.     print(12345) 
14.     print('xxxx') 
15. print('退出縮進(jìn)') 

運(yùn)行效果如下圖所示：

可以看出以下特點(diǎn)：

1. 被添加的回調(diào)函數(shù)進(jìn)入了一個棧，所以后添加的回調(diào)函數(shù)先調(diào)用
2. 回調(diào)函數(shù)會在結(jié)束縮進(jìn)的時候被調(diào)用

現(xiàn)在我們來人工構(gòu)造一個異常：

可以看到，即使縮進(jìn)里面出現(xiàn)了報錯，回調(diào)函數(shù)仍然可以正常運(yùn)行。等所有回調(diào)函數(shù)運(yùn)行完成以后，Python 才會退出。

基于以上特點(diǎn)，我們就可以來重構(gòu)最開始的代碼了：

1. import json 
2. import redis 
3. import datetime 
4. import pymongo 
5. import contextlib 
6.  
7. client = redis.Redis() 
8. handler = pymongo.MongoClient().test.data 
9.  
10. def add_ts(): 
11.     handler.insert_one({'finished': True, 'ts': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}) 
12.  
13. def parse(): 
14.     with contextlib.ExitStack() as stack: 
15.         stack.callback(add_ts) 
16.         while True: 
17.             data = client.lpop('test') 
18.             if not data: 
19.                 return 
20.             handler.insert_one(json.loads(data)) 

無論是正常運(yùn)行結(jié)束還是運(yùn)行過程中報錯，add_ts函數(shù)都會正常運(yùn)行，確保始終增加一條日期數(shù)據(jù)。

原文鏈接：https://mp.weixin.qq.com/s/fjbOA0YFW1jrsthtNem5aQ