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 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关键字执行,形如:

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

的问题。。。

TIL

总之学到了contextlib里的一些有用的函数和装饰器,也第一次发现with可以放多个context

虽然放多个context的动态构造还有待研究,with 后面的代码块也不能填一个元组或者列表。。惆怅。。

文章被以下专栏收录

    微信公众号同名,欢迎投稿。telegram:t.me/PyChina,由@阿橙FM 发起召集,全平台约20万开发者关注,会员来自全球十多个国家和地区,拥有十多个线上线下技术社群,向本专栏投稿即默认发布到Python中文社区全平台。GitHub:github.com/pycn