python with提前退出遇到的坑与解决方案_Python

问题的起源

早些时候使用with实现了一版全局进程锁，希望实现以下效果：

									with CacheLock("test_lock", 10):

									  #如果抢占到锁，那么就执行这段代码

									  # 否则，让with提前退出

全局进程锁本身不用多说，大部分都依靠外部的缓存来实现的，redis上用的是setnx，有时候根据需要加上缓存击穿问题、随机延后以防止对缓存本身造成压力

当时同样写了单元测试来测试这段代码的有效性：

									with CacheLock("test_lock", 10):

									  value = cache.get("test_lock")

									  self.assertEqual(value, 1)

									  with CacheLock("test_lock", 10):

									    # 不会进到这里

									    self.assertFalse(True)

									value = cache.get("test_lock")

									self.assertEqual(value, None)

看起来非常完美地通过了。

这样的一个全局进程锁是通过 __enter__ 方法抛出异常, __exit__ 方法中捕获异常来实现的:

									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

看起来还不错，毕竟单元测试都过了。

但是，这样的实现是有问题的：

原因在于 __exit__ 的执行不是包在 __enter__ 之外的，因此 __enter__ 抛出的异常，不会被 __exit__ 捕获。

上面的单元测试恰好通过，是因为其中有两个with语句，外面的with 捕获的其实是里面的 __enter__ 抛出的异常

使用改进后的单元测试：

									cache.set("test_lock",1)

									with CacheLock("test_lock", 10):

									  self.assertFalse(True)

									value = cache.get("test_lock")

									self.assertEqual(value, None)

就会发现单元测试过不去了。

这个问题是我试图使用with实现另一个逻辑：AB测试时出现的，同样是 __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):

									    # 如果不满足条件，等于不执行上下文中的内容

									    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

调试没有通过的单元测试的时候发现，抛出异常后根本没有执行到 __enter__

第一种解决方案

既然想明白了with的执行顺序，那么第一种解决方案就呼之欲出了：既然__exit__捕获的异常在__enter__执行完成之后，那么我们提供一个函数确认一下就可以了，把ABContext实现改成这样：

									def ensure(self):

									    if not self.relationship(self.version, self.ab_layer.value):

									      raise EarlyExit("not match")

									  def __enter__(self):

									    # 如果不满足条件，等于不执行上下文中的内容

									    return self

使用的时候：

									with ABContext(newVersion, consts.ABEnum.layer2) as c:

									  c.ensure()

									  # 执行其他的想要执行的代码

但这样的解决方法并不优雅，万一使用这个ABContext的时候忘记用ensure方法了，那么就等于完全没用这个Context方法，太容易失误了，而且代码也失去了Pythonic的性质

第二种解决方法

翻了一下contextlib的标准库文档，发现有一个已经废弃的函数： contextlib.nested

									from contextlib import nested

									with nested(*managers):

									  do_something()

可以执行多个上下文.

									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关键字执行，形如：

1 2	`with context1,context2:` `#do something`

这个特性还不错，根据 __enter__ 的执行顺序的话，那么我们可以实现一个由第一个 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

结合前面我们实现的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，单元测试就这样过了

能不能再给力点？

确实，在with里要写俩context有点蛋疼，并不是特别优雅，能不能还是回到最初的那种用法：我们只用写一条context，这一个context做到了两个context的事情？

要是nested那个函数还在就好了。。要的其实就是它的功能。

Python3.1之后contextlib提供了一个ExitStack的功能来提供一个模拟的功能，但试了一下发现，实际上只调用了__enter__方法，但没有做对应的异常捕获

第三种解决方案

哈哈哈哈把自己绕到圈子里去了，想了一下，同样是一个缩进的代码块，为什么不能用if来解决呢！不就是个：

									def test_context_noteq(self):

									    # 不等的时候，不会执行with里的内容

									    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()