python with提前退出：坑與解決方案

問題的起源

早些時候使用with實現了一版全局進程鎖，希望實現以下效果：

with CacheLock("test_lock", 10):n #如果搶佔到鎖，那麼就執行這段代碼n # 否則，讓with提前退出n

全局進程鎖本身不用多說，大部分都依靠外部的緩存來實現的，redis上用的是setnx，有時候根據需要加上緩存擊穿問題、隨機延後以防止對緩存本身造成壓力

當時同樣寫了單元測試來測試這段代碼的有效性：

with CacheLock("test_lock", 10):n value = cache.get("test_lock")n self.assertEqual(value, 1)n with CacheLock("test_lock", 10):n # 不會進到這裡n self.assertFalse(True)nvalue = cache.get("test_lock")nself.assertEqual(value, None)n

看起來非常完美地通過了。

這樣的一個全局進程鎖是通過__enter__方法拋出異常, __exit__方法中捕獲異常來實現的:

class CacheLock(object):n def __init__(self, lock_key, lock_timeout):n self.lock_key = lock_keyn self.lock_timeout = lock_timeoutn self.success = Falsen def __enter__(self):n self.success = cache.lock(self.lock_key, self.lock_timeout)n if self.success:n return selfn else:n raise LockException("not have lock")nn def __exit__(self, exc_type, exc_value, traceback):n #沒有搶到鎖的時候，啥都不做？n if self.success:n await cache.delete(self.lock_key)n if isinstance(exc_value, LockException):n return Truen if exc_type:n raise exc_valuen

看起來還不錯，畢竟單元測試都過了。

但是，這樣的實現是有問題的：

原因在於__exit__ 的執行不是包在__enter__ 之外的，因此__enter__拋出的異常，不會被__exit__捕獲。

上面的單元測試恰好通過，是因為其中有兩個with語句，外面的with 捕獲的其實是裡面的__enter__ 拋出的異常

使用改進後的單元測試：

cache.set("test_lock",1)nwith CacheLock("test_lock", 10):n self.assertFalse(True)nvalue = cache.get("test_lock")nself.assertEqual(value, None)n

就會發現單元測試過不去了。

這個問題是我試圖使用with實現另一個邏輯：AB測試 時出現的，同樣是__enter__拋出異常，__exit__ 試圖捕獲：

import operatorn nclass EarlyExit(Exception):n passnnclass ABContext(object):n """AB測試上下文n >>> with ABContext(newVersion, consts.ABEnum.layer2):n >>> # dosomethingn """nn def __init__(self, version, ab_layer, relationship="eq"):n self.version = versionn self.ab_layer = ab_layern # 如果不存在這種操作符，那就提前報錯n self.relationship = getattr(operator, relationship)nn def __enter__(self):n # 如果不滿足條件，等於不執行上下文中的內容n if not self.relationship(self.version, self.ab_layer.value):n raise EarlyExit("not match")n return selfnn def __exit__(self, exc_type, exc_value, traceback):n if exc_value is None:n return Truen if isinstance(exc_value, EarlyExit):n return Truen if exc_type:n raise exc_valuen return Truen

調試沒有通過的單元測試的時候發現，拋出異常後根本沒有執行到__enter__

第一種解決方案

既然想明白了with的執行順序，那麼第一種解決方案就呼之欲出了：既然__exit__捕獲的異常在__enter__執行完成之後，那麼我們提供一個函數確認一下就可以了，把ABContext實現改成這樣：

def ensure(self):n if not self.relationship(self.version, self.ab_layer.value):n raise EarlyExit("not match")nn def __enter__(self):n # 如果不滿足條件，等於不執行上下文中的內容n return selfn

使用的時候：

with ABContext(newVersion, consts.ABEnum.layer2) as c:n c.ensure()n # 執行其他的想要執行的代碼n

但這樣的解決方法並不優雅，萬一使用這個ABContext的時候忘記用ensure方法了，那麼就等於完全沒用這個Context方法，太容易失誤了，而且代碼也失去了Pythonic的性質

第二種解決方法

翻了一下contextlib的標準庫文檔，發現有一個已經廢棄的函數：contextlib.nested

from contextlib import nestednnwith nested(*managers):n do_something() n

可以執行多個上下文.

from contextlib import nestednwith nested(A(), B(), C()) as (X, Y, Z):n do_something()n# is equivalent to this:nm1, m2, m3 = A(), B(), C()nwith m1 as X:n with m2 as Y:n with m3 as Z:n do_something()n

這個廢棄的特性在Python2.7之後，可以直接由with關鍵字執行，形如：

with context1,context2:n #do somethingn

這個特性還不錯，根據__enter__的執行順序的話，那麼我們可以實現一個由第一個 context的__exit__來捕獲，第二個context的__enter__來拋出異常，

如同這樣：

class AlwaySuccessContext(object):n def __enter__(self):n return selfnn def __exit__(self, exc_type, exc_value, traceback):n if isinstance(exc_value, EarlyExit):n return Truen if exc_type:n raise exc_valuen return Truen

結合前面我們實現的ABContext的使用是這樣的：

def test_context_noteq(self):n obj = MagicMock(return_value=True)n with AlwaySuccessContext(), ABContext(2, const.ABTestEnum.control):n self.assertFalse(obj())n obj.assert_not_called()n

good，單元測試就這樣過了

能不能再給力點？

確實，在with里要寫倆context有點蛋疼，並不是特別優雅，能不能還是回到最初的那種用法：我們只用寫一條context，這一個context做到了兩個context的事情？

要是nested那個函數還在就好了。。要的其實就是它的功能。

Python3.1之後contextlib提供了一個ExitStack的功能來提供一個模擬的功能，但試了一下發現，實際上只調用了__enter__方法，但沒有做對應的異常捕獲

第三種解決方案

哈哈哈哈把自己繞到圈子裡去了，想了一下，同樣是一個縮進的代碼塊，為什麼不能用if來解決呢！不就是個：

def test_context_noteq(self):n # 不等的時候，不會執行with里的內容n obj = MagicMock(return_value=True)n context = ABContext(2, const.ABTestEnum.control)n # print(type(context))n if ABContext(2, const.ABTestEnum.control):n self.assertFalse(obj())n obj.assert_not_called()n

的問題。。。

TIL

總之學到了contextlib里的一些有用的函數和裝飾器，也第一次發現with可以放多個context

雖然放多個context的動態構造還有待研究，with 後面的代碼塊也不能填一個元組或者列表。。惆悵。。