python-contextlib上下文管理器

发布时间 2023-07-05 09:31:28作者: 贝壳里的星海

python contextlib上下文管理器

python-contextlib

  • 上下文管理器 两大作用:
    -- 可以以一种更加优雅的方式,操作(创建/获取/释放)资源,如文件操作、数据库连接;
    -- 可以以一种更加优雅的方式,处理异常

读取文件的一般流程

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

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

使用with实现

with open('filename.csv') as file:
    print(file.readlines())
with context_expression [as target(s)]:
    with-body

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

类的内部实现

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

  • __enter__:在进入 with 语法块之前调用,返回值会赋值给 withtarget
  • __exit__:在退出 with 语法块时调用,一般用作异常处理, 除了 self 之外,必须传入另外三个参数,分别表示 exception 的类型
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

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

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

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

  • exc_type:异常类型
  • exc_value:异常对象
  • exc_tb:异常堆栈信息
class Sample():
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):

        if exc_type == IndexError:
            print("执行出错 IndexError")
            print(exc_value, type(exc_value))
            print(traceback)
            return True
        elif exc_type == ZeroDivisionError:
            print("执行出错 ZeroDivisionError")
            print(exc_value, type(exc_value))
            print(traceback)
            return True

    def do_wrong_job(self):
        a = 1 / 0
        # 1 / 0 ,是一个错误的式子,应该会报错
        return a


with Sample() as sample:
    sample.do_wrong_job()

    
# 执行结果
执行出错 ZeroDivisionError
division by zero <class 'ZeroDivisionError'>
<traceback object at 0x00000237CF4F59C0>

contextlib模块

对于需要上下文管理的场景,除了自己实现 __enter____exit__ 之外,还有更简单的方式,只写一个函数就可以实现上下文管理器

contextlib 是一个装饰器,只要按照它的代码协议来实现函数内容,就可以将这个函数对象变成一个上下文管理器。

__all__ = ["asynccontextmanager", "contextmanager", "closing", "nullcontext",
           "AbstractContextManager", "AbstractAsyncContextManager",
           "AsyncExitStack", "ContextDecorator", "ExitStack",
           "redirect_stdout", "redirect_stderr", "suppress"]

contextmanager装饰器

生成器函数转为上下文管理器

@contextmanagercontextlib 模块提供的一个装饰器,可以将一个生成器函数转换成一个上下文管理器

from contextlib import contextmanager

@contextmanager
def my_context():
    # 进入上下文前的操作
    print('entering context')
    try:
        yield  # 生成器函数的 `yield` 语句之前的代码作为上文,之后的代码作为下文
    finally:
        # 离开上下文后的操作
        print('exiting context')

# 使用上下文管理器
with my_context():
    print('inside context')

定义了一个 my_context 上下文管理器,使用 @contextmanager 装饰器将其转换为一个生成器函数。在该生成器函数中,我们可以在 yield 语句前后添加上下文的进入和离开时的操作。在使用上下文管理器时,可以使用 with 语句来自动管理上下文,而不需要手动调用 __enter____exit__ 方法。

当执行 with my_context() 语句时,程序会进入 my_context 上下文,执行 entering context,然后执行 yield 语句前的代码,即 inside context。当程序从 with 语句块中退出时,会自动执行 finally 语句块中的代码,即 exiting context

从生成器到上下文管理器

@contextmanager 装饰器可以帮助我们更轻松地创建上下文管理器,避免手动编写 __enter____exit__ 方法的繁琐。

需要主要的是这个函数必须是个装饰器
被装饰器装饰的函数分为三部分:
with 语句中的代码块执行前执行函数中 yield 之前代码
yield 返回的内容复制给 as 之后的变量
with 代码块执行完毕后执行函数中 yield 之后的代码

再简单理解就是
yield 前半段用来表示__enter__()
yield 后半段用来表示__exit__()

from contextlib import contextmanager

@contextmanager
def open_func(file_name):
    # __enter__方法
    print('open file:', file_name)
    file_handler = open(file_name, 'r')

    yield file_handler  # 这里的 yield 用于返回一个生成器

    # __exit__方法
    print('close file:', file_name)
    file_handler.close()
    return


with open_func(r'filename.txt') as file_in:
    for line in file_in:
        print(line)
        

# 执行结果        
open file: filename.txt
12344567890
close file: filename.txt

使用继承类管理上下文

ContextDecorator

contextlib.ContextDecorator 是一个抽象基类,通过继承 contextlib 里面的 ContextDecorator 类,用来定义上下文管理器装饰器。使用它可以使得上下文管理器可以像装饰器一样使用。

我们可以通过继承 ContextDecorator 类来创建一个上下文管理器装饰器。

from contextlib import ContextDecorator

class my_decorator(ContextDecorator):
    def __enter__(self):
        print('Entering the context')
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        print('Exiting the context')
        return False

@my_decorator()
def my_function():
    print('Inside the function')

my_function()
# 直接结果
Entering the context
Inside the function
Exiting the context

关闭open的句柄

并不是所有的类都支持上下文管理器的 API,有一些遗留的类会使用一个 close 方法。
为了确保关闭句柄,需要使用 closing 为他创建一个上文管理器。

import contextlib

class Http():
    def __init__(self):
        self.session = "open"
        print("init: set open")

    def close(self):
        """
        关闭的方法必须叫 close
        """
        self.session = "close"
        print("set close")


if __name__ == '__main__':
    with contextlib.closing(Http()) as http:
        print(f"inside session value:{http.session}")
    print(f"outside session value:{http.session}")

    with contextlib.closing(Http()) as http:
        print(f"inside session value:{http.session}")
        raise EnvironmentError("EnvironmentError")
    print(f"outside session value:{http.session}")
    
    
#  直接结果
init: set open
inside session value:open
set close
outside session value:close
init: set open
inside session value:open
set close
Traceback (most recent call last):
  File "D:/Note/lcodeNoteCards/testcode/python/testpy.py", line 24, in <module>
    raise EnvironmentError("EnvironmentError")
OSError: EnvironmentError

即使程序出现了错误,最后也会执行 close 方法的内容。

参考资料

参考资料1