Python进阶——什么是上下文管理器?

·  阅读 1072
Python进阶——什么是上下文管理器?

关注「水滴与银弹」公众号,第一时间获取优质技术干货。7年资深后端研发,用简单的方式把技术讲清楚。

阅读本文大约需要 12 分钟。

在 Python 开发中,我们经常会使用到 with 语法块,例如在读写文件时,保证文件描述符的正确关闭,避免资源泄露问题。

你有没有思考过, with 背后是如何实现的?我们常常听到的上下文管理器究竟是什么?

这篇文章我们就来学习一下 Python 上下文管理器,以及 with 的运行原理。

with语法块

在讲解 with 语法之前,我们先来看一下不使用 with 的代码如何写?

我们在操作一个文件时,代码可以这么写:

# 打开文件
f = open('file.txt')
for line in f:
    # 读取文件内容 执行其他操作
    # do_something...
# 关闭文件
f.close()
复制代码

这个例子非常简单,就是打开一个文件,然后读取文件中的内容,最后关闭文件释放资源。

但是,代码这么写会有一个问题:在打开文件后,如果要对读取到的内容进行其他操作,在这操作期间发生了异常,这就会导致文件句柄无法被释放,进而导致资源的泄露。

如何解决这个问题?

也很简单,我们使用 try ... finally 来优化代码:

# 打开文件
f = open('file.txt')
try:
    for line in f:
        # 读取文件内容 执行其他操作
        # do_something...
finally:
    # 保证关闭文件
    f.close()
复制代码

这么写的好处是,在读取文件内容和操作期间,无论是否发生异常,都可以保证最后能释放文件资源。

但这么优化,代码结构会变得很繁琐,每次都要给代码逻辑增加 try ... finally 才可以,可读性变得很差。

针对这种情况,我们就可以使用 with 语法块来解决这个问题:

with open('file.txt') as f:
    for line in f:
        # do_something...
复制代码

使用 with 语法块可以完成之前相同的功能,而且这么写的好处是,代码结构变得非常清晰,可读性也很好。

明白了 with 的作用,那么 with 究竟是如何运行的呢?

上下文管理器

首先,我们来看一下 with 的语法格式:

with context_expression [as target(s)]:
    with-body
复制代码

with 语法非常简单,我们只需要 with 一个表达式,然后就可以执行自定义的业务逻辑。

但是,with 后面的表达式是可以任意写的吗?

答案是否定的。要想使用 with 语法块,with 后面的的对象需要实现「上下文管理器协议」。

什么是「上下文管理器协议」?

一个类在 Python 中,只要实现以下方法,就实现了「上下文管理器协议」:

  • __enter__:在进入 with 语法块之前调用,返回值会赋值给 withtarget
  • __exit__:在退出 with 语法块时调用,一般用作异常处理

我们来看实现了这 2 个方法的例子:

class TestContext:

    def __enter__(self):
        print('__enter__')
        return 1

    def __exit__(self, exc_type, exc_value, exc_tb):
        print('exc_type: %s' % exc_type)
        print('exc_value: %s' % exc_value)
        print('exc_tb: %s' % exc_tb)

with TestContext() as t:
    print('t: %s' % t)
    
# Output:
# __enter__
# t: 1
# exc_type: None
# exc_value: None
# exc_tb: None
复制代码

在这个例子中,我们定义了 TestContext 类,它分别实现了 __enter____exit__ 方法。

这样一来,我们就可以把 TestContext 当做一个「上下文管理器」来使用,也就是通过 with TestContext() as t 方式来执行。

从输出结果我们可以看到,具体的执行流程如下:

  • __enter__ 在进入 with 语句块之前被调用,这个方法的返回值赋给了 with 后的 t 变量
  • __exit__ 在执行完 with 语句块之后被调用

如果在 with 语句块内发生了异常,那么 __exit__ 方法可以拿到关于异常的详细信息:

  • exc_type:异常类型
  • exc_value:异常对象
  • exc_tb:异常堆栈信息

我们来看一个发生异常的例子,观察 __exit__ 方法拿到的异常信息是怎样的:

with TestContext() as t:
    # 这里会发生异常
    a = 1 / 0 
    print('t: %s' % t)

# Output:
# __enter__
# exc_type: <type 'exceptions.ZeroDivisionError'>
# exc_value: integer division or modulo by zero
# exc_tb: <traceback object at 0x10d66dd88>
# Traceback (most recent call last):
#   File "base.py", line 16, in <module>
#     a = 1 / 0
# ZeroDivisionError: integer division or modulo by zero
复制代码

从输出结果我们可以看到,当 with 语法块内发生异常后,__exit__ 输出了这个异常的详细信息,其中包括异常类型、异常对象、异常堆栈。

如果我们需要对异常做特殊处理,就可以在这个方法中实现自定义逻辑。

回到最开始我们讲的,使用 with 读取文件的例子。之所以 with 能够自动关闭文件资源,就是因为内置的文件对象实现了「上下文管理器协议」,这个文件对象的 __enter__ 方法返回了文件句柄,并且在 __exit__ 中实现了文件资源的关闭,另外,当 with 语法块内有异常发生时,会抛出异常给调用者。

伪代码可以这么写:

class File:

    def __enter__(self):
        return file_obj

    def __exit__(self, exc_type, exc_value, exc_tb):
        # with 退出时释放文件资源
        file_obj.close()
        # 如果 with 内有异常发生 抛出异常
        if exc_type is not None:
            raise exception
复制代码

这里我们小结一下,通过对 with 的学习,我们了解到,with 非常适合用需要对于上下文处理的场景,例如操作文件、Socket,这些场景都需要在执行完业务逻辑后,释放资源。

contextlib模块

对于需要上下文管理的场景,除了自己实现 __enter____exit__ 之外,还有更简单的方式来做吗?

答案是肯定的。我们可以使用 Python 标准库提供的 contextlib 模块,来简化我们的代码。

使用 contextlib 模块,我们可以把上下文管理器当成一个「装饰器」来使用。

其中,contextlib 模块提供了 contextmanager 装饰器和 closing 方法。

下面我们通过例子来看一下它们是如何使用的。

contextmanager装饰器

我们先来看 contextmanager 装饰器的使用:

from contextlib import contextmanager

@contextmanager
def test():
    print('before')
    yield 'hello'
    print('after')

with test() as t:
    print(t)

# Output:
# before
# hello
# after
复制代码

在这个例子中,我们使用 contextmanager 装饰器和 yield配合,实现了和前面上下文管理器相同的功能,它的执行流程如下:

  1. 执行 test() 方法,先打印出 before
  2. 执行 yield 'hello'test 方法返回,hello 返回值会赋值给 with 语句块的 t 变量
  3. 执行 with 语句块内的逻辑,打印出 t 的值 hello
  4. 又回到 test 方法中,执行 yield 后面的逻辑,打印出 after

这样一来,当我们使用这个 contextmanager 装饰器后,就不用再写一个类来实现上下文管理协议,只需要用一个方法装饰对应的方法,就可以实现相同的功能。

不过有一点需要我们注意:在使用 contextmanager 装饰器时,如果被装饰的方法内发生了异常,那么我们需要在自己的方法中进行异常处理,否则将不会执行 yield 之后的逻辑。

@contextmanager
def test():
    print('before')
    try:
        yield 'hello'
        # 这里发生异常 必须自己处理异常逻辑 否则不会向下执行
        a = 1 / 0 
    finally:
        print('after')

with test() as t:
    print(t)
复制代码

closing方法

我们再来看 contextlib 提供的 closing 方法如何使用。

closing 主要用在已经实现 close 方法的资源对象上:

from contextlib import closing

class Test():

    # 定义了 close 方法才可以使用 closing 装饰器
    def close(self):
        print('closed')

# with 块执行结束后 自动执行 close 方法
with closing(Test()):
    print('do something')
    
# Output:
# do something
# closed
复制代码

从执行结果我们可以看到,with 语句块执行结束后,会自动调用 Test 实例的 close 方法。

所以,对于需要自定义关闭资源的场景,我们可以使用这个方法配合 with 来完成。

contextlib的实现

学习完了 contextlib 模块的使用,最后我们来看一下 contextlib 模块是究竟是如何实现的?

contextlib 模块相关的源码如下:

class _GeneratorContextManagerBase:

    def __init__(self, func, args, kwds):
        # 接收一个生成器对象 (方法内包含 yield 的方法就是一个生成器)
        self.gen = func(*args, **kwds)
        self.func, self.args, self.kwds = func, args, kwds
        doc = getattr(func, "__doc__", None)
        if doc is None:
            doc = type(self).__doc__
        self.__doc__ = doc

class _GeneratorContextManager(_GeneratorContextManagerBase,
                               AbstractContextManager,
                               ContextDecorator):

    def __enter__(self):
        try:
            # 执行生成器 代码会运行生成器方法的 yield 处
            return next(self.gen)
        except StopIteration:
            raise RuntimeError("generator didn't yield") from None

    def __exit__(self, type, value, traceback):
        # with 内没有异常发生
        if type is None:
            try:
                # 继续执行生成器
                next(self.gen)
            except StopIteration:
                return False
            else:
                raise RuntimeError("generator didn't stop")
        # with 内发生了异常
        else:
            if value is None:
                value = type()
            try:
                # 抛出异常
                self.gen.throw(type, value, traceback)
            except StopIteration as exc:
                return exc is not value
            except RuntimeError as exc:
                if exc is value:
                    return False
                if type is StopIteration and exc.__cause__ is value:
                    return False
                raise
            except:
                if sys.exc_info()[1] is value:
                    return False
                raise
            raise RuntimeError("generator didn't stop after throw()")

def contextmanager(func):
    @wraps(func)
    def helper(*args, **kwds):
        return _GeneratorContextManager(func, args, kwds)
    return helper

class closing(AbstractContextManager):
    def __init__(self, thing):
        self.thing = thing
    def __enter__(self):
        return self.thing
    def __exit__(self, *exc_info):
        self.thing.close()
复制代码

源码中我已经添加好了注释,你可以详细看一下。

contextlib 源码中逻辑其实比较简单,其中 contextmanager 装饰器实现逻辑如下:

  1. 初始化一个 _GeneratorContextManager 类,构造方法接受了一个生成器 gen
  2. 这个类实现了上下文管理器协议 __enter____exit__
  3. 执行 with 时会进入到 __enter__ 方法,然后执行这个生成器,执行时会运行到 with 语法块内的 yield
  4. __enter__ 返回 yield 的结果
  5. 如果 with 语法块没有发生异常,with 执行结束后,会进入到 __exit__ 方法,再次执行生成器,这时会运行 yield 之后的代码逻辑
  6. 如果 with 语法块发生了异常,__exit__ 会把这个异常通过生成器,传入到 with 语法块内,也就是把异常抛给调用者

再来看 closing 的实现,closing 方法就是在 __exit__ 方法中调用了自定义对象的 close,这样当 with 结束后就会执行我们定义的 close 方法。

使用场景

学习完了上下文管理器,那么它们具体会用在什么场景呢?

下面我举几个常用的例子来演示下,你可以参考一下结合自己的场景使用。

Redis分布式锁

from contextlib import contextmanager

@contextmanager
def lock(redis, lock_key, expire):
    try:
        locked = redis.set(lock_key, 'locked', expire)
        yield locked
    finally:
        redis.delete(lock_key)

# 业务调用 with 代码块执行结束后 自动释放锁资源
with lock(redis, 'locked', 3) as locked:
    if not locked:
        return
    # do something ...
复制代码

在这个例子中,我们实现了 lock 方法,用于在 Redis 上申请一个分布式锁,然后使用 contextmanager 装饰器装饰了这个方法。

之后我们业务在调用 lock 方法时,就可以使用 with 语法块了。

with 语法块的第一步,首先判断是否申请到了分布式锁,如果申请失败,则业务逻辑直接返回。如果申请成功,则执行具体的业务逻辑,当业务逻辑执行完成后,with 退出时会自动释放分布式锁,就不需要我们每次都手动释放锁了。

Redis事物和管道

from contextlib import contextmanager

@contextmanager
def pipeline(redis):
    pipe = redis.pipeline()
    try:
        yield pipe
        pipe.execute()
    except Exception as exc:
        pipe.reset()
            
# 业务调用 with 代码块执行结束后 自动执行 execute 方法
with pipeline(redis) as pipe:
    pipe.set('key1', 'a', 30)
    pipe.zadd('key2', 'a', 1)
    pipe.sadd('key3', 'a')
复制代码

在这个例子中,我们定义了 pipeline 方法,并使用装饰器 contextmanager 让它变成了一个上下文管理器。

之后在调用 with pipeline(redis) as pipe 时,就可以开启一个事物和管道,然后在 with 语法块内向这个管道中添加命令,最后 with 退出时会自动执行 pipelineexecute 方法,把这些命令批量发送给 Redis 服务端。

如果在执行命令时发生了异常,则会自动调用 pipelinereset 方法,放弃这个事物的执行。

总结

总结一下,这篇文章我们主要介绍了 Python 上下文管理器的使用及实现。

首先我们介绍了不使用 with 和使用 with 操作文件的代码差异,然后了解到使用 with 可以让我们的代码结构更加简洁。之后我们探究了 with 的实现原理,只要实现 __enter____exit__ 方法的实例,就可以配合 with 语法块来使用。

之后我们介绍了 Python 标准库的 contextlib 模块,它提供了实现上下文管理更好的使用方式,我们可以使用 contextmanager 装饰器和 closing 方法来操作我们的资源。

最后我举了两个例子,来演示上下文管理器的具体使用场景,例如在 Redis 中使用分布式锁和事物管道,用上下文管理器帮我们管理资源,执行前置和后置逻辑。

所以,如果我们在开发中把操作资源的前置和后置逻辑,通过上下文管理器来实现,那么我们的代码结构和可维护性也会有所提高,推荐使用起来。

我的 Python 进阶系列文章:

爬虫系列文章:

想看更多硬核技术文章?关注「水滴与银弹」公众号,第一时间获取优质技术干货。7年资深后端研发,用简单的方式把技术讲清楚。

分类:
后端
标签:
收藏成功!
已添加到「」, 点击更改