快捷導(dǎo)航

python with提前退出遇到的坑與解決方案

更新時間：2018年01月05日 11:01:11 作者：豬了個去

這篇文章主要介紹了python with提前退出遇到的坑與解決方法，需要的朋友參考下吧

問題的起源

早些時候使用with實現(xiàn)了一版全局進(jìn)程鎖，希望實現(xiàn)以下效果：

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

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

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

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

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

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

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

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

但是，這樣的實現(xiàn)是有問題的：

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

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

使用改進(jìn)后的單元測試：

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

就會發(fā)現(xiàn)單元測試過不去了。

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

import operator
class EarlyExit(Exception):
  pass
class ABContext(object):
  """AB測試上下文
  >>> with ABContext(newVersion, consts.ABEnum.layer2):
  >>>   # dosomething
  """
  def __init__(self, version, ab_layer, relationship="eq"):
    self.version = version
    self.ab_layer = ab_layer
    # 如果不存在這種操作符，那就提前報錯
    self.relationship = getattr(operator, relationship)
  def __enter__(self):
    # 如果不滿足條件，等于不執(zhí)行上下文中的內(nèi)容
    if not self.relationship(self.version, self.ab_layer.value):
      raise EarlyExit("not match")
    return self
  def __exit__(self, exc_type, exc_value, traceback):
    if exc_value is None:
      return True
    if isinstance(exc_value, EarlyExit):
      return True
    if exc_type:
      raise exc_value
    return True

調(diào)試沒有通過的單元測試的時候發(fā)現(xiàn)，拋出異常后根本沒有執(zhí)行到 __enter__

第一種解決方案

既然想明白了with的執(zhí)行順序，那么第一種解決方案就呼之欲出了：既然__exit__捕獲的異常在__enter__執(zhí)行完成之后，那么我們提供一個函數(shù)確認(rèn)一下就可以了，把ABContext實現(xiàn)改成這樣：

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

使用的時候：

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

但這樣的解決方法并不優(yōu)雅，萬一使用這個ABContext的時候忘記用ensure方法了，那么就等于完全沒用這個Context方法，太容易失誤了，而且代碼也失去了Pythonic的性質(zhì)

第二種解決方法

翻了一下contextlib的標(biāo)準(zhǔn)庫文檔，發(fā)現(xiàn)有一個已經(jīng)廢棄的函數(shù)： contextlib.nested

from contextlib import nested
with nested(*managers):
  do_something()

可以執(zhí)行多個上下文.

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

這個廢棄的特性在Python2.7之后，可以直接由with關(guān)鍵字執(zhí)行，形如：

with context1,context2:
  #do something

這個特性還不錯，根據(jù) __enter__ 的執(zhí)行順序的話，那么我們可以實現(xiàn)一個由第一個 context的__exit__來捕獲，第二個context的__enter__來拋出異常，

如同這樣：

class AlwaySuccessContext(object):
  def __enter__(self):
    return self
  def __exit__(self, exc_type, exc_value, traceback):
    if isinstance(exc_value, EarlyExit):
      return True
    if exc_type:
      raise exc_value
    return True

結(jié)合前面我們實現(xiàn)的ABContext的使用是這樣的：

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

good，單元測試就這樣過了

能不能再給力點(diǎn)？

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

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

Python3.1之后contextlib提供了一個ExitStack的功能來提供一個模擬的功能，但試了一下發(fā)現(xiàn)，實際上只調(diào)用了__enter__方法，但沒有做對應(yīng)的異常捕獲

第三種解決方案

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

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

TIL

總之學(xué)到了contextlib里的一些有用的函數(shù)和裝飾器，也第一次發(fā)現(xiàn)with可以放個context

雖然放多個context的動態(tài)構(gòu)造還有待研究，with 后面的代碼塊也不能填一個元組或者列表。。惆悵。。

您可能感興趣的文章: