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

问题的起源

早些时候使用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 operatorclass EarlyExit(Exception):  passclass 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