Python 上下文管理器

测试有道

于 2024-08-16 12:38:47 发布

阅读量761

点赞数 20

分类专栏： Python 文章标签： python 开发语言

本文链接：https://blog.csdn.net/lIujunXHU/article/details/141255685

版权

Python 专栏收录该内容

13 篇文章 1 订阅

订阅专栏

1.4.1、contextmanager 装饰器

1、上下文管理器

1.1、为什么需要上下文管理器？

在日常编码中我们经常会操作一些资源，比如文件对象、数据库连接、Socket连接等，资源操作完了之后，不管操作的成功与否，都要关闭该资源，避免出现内存泄漏，CPU资源占用的问题。如果没有上下文管理器，我们关闭资源就需要手动关闭，但是这里可能会被遗漏，所以我们就需要一个自动帮我们关闭资源的东西，那就是上下文管理器。
示例：手动关闭资源的场景（手动关闭文件）

try:
    # 打开文件
    f = open("XXX", "w", encoding='utf-8')
    f.write("hello world")
except Exception as e:
    # 处理操作文件错误
    print(f"操作文件错误，错误原因：{e!r}")
finally:
    # 关闭文件
    f.close()

1.2、上下文管理器的原理

1.2.1、实现协议

上下文管理器协议：是指要实现对象的 __enter__() 和 __exit__() 方法。

__enter__(self)：进入上下文管理器的运行时上下文。在with语句块执行之前调用，其返回值将绑定到with语句的as子句指定的变量上。
__exit__(self, exc_type, exc_val, exc_tb)：退出上下文管理器的运行时上下文。在with语句块执行完毕后调用，无论是正常结束还是异常结束。exc_type、exc_val和exc_tb分别表示异常类型、异常值和traceback信息，如果with语句块没有异常发生，则这三个参数都将为None。

示例：自定义实现一个上下文管理器

class MyContextManager:  
    def __enter__(self):  
        print("Enter the context")  
        return self  
      
    def __exit__(self, exc_type, exc_val, exc_tb):  
        print("Exit the context")  
        if exc_type is not None:  
            print(f"Error occurred: {exc_type}")  
        return False  # 如果为True，会阻止异常的抛出  
  
with MyContextManager() as cm:  
    print("Inside the context")  
    # raise Exception("Something bad happened")  # 你可以尝试取消注释这行代码，看看`__exit__`如何处理异常

1.2.2、异常处理

1、如果上下文是无异常地退出的，三个参数都将为None。
2、如果 with_body中出现了异常，并且希望方法屏蔽此异常（即避免其被传播），则应当返回真值即返回True，没有显示返回True默认返回None（False）。否则的话，异常将在退出方法__exit__时按正常流程处理。
3、请注意__exit__()方法不应该重新引发被传入的异常，这是调用者的责任。如果 with_body 的退出由异常引发，并且__exit__()的返回值等于 False，那么这个异常将被重新引发一次；如果 __exit__() 的返回值等于 True，此时认为异常在__exit__中以及处理过了，那么这个异常就被无视掉，继续执行后面的代码。

示例1：__exit__() 返回True，无视异常，继续执行后续代码

class MyContextManager:
    def __enter__(self):
        print("Enter the context")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Exit the context")
        if exc_type is not None:
            print(f"Error occurred: {exc_type}")
        return True  # 如果为True，会阻止异常的抛出


with MyContextManager() as cm:
    print("Inside the context")
    raise Exception("Something bad happened")

示例2：__exit__() 返回False，抛出异常，中断执行代码

class MyContextManager:
    def __enter__(self):
        print("Enter the context")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Exit the context")
        if exc_type is not None:
            print(f"Error occurred: {exc_type}")
        return False  # 如果为True，会阻止异常的抛出


with MyContextManager() as cm:
    print("Inside the context")
    raise Exception("Something bad happened")

1.2.3、with语句执行原理

上下文管理器with语法：

with context_obj [as obj]:
    with_body

执行原理：

1、使用with语句使用的时，上下文管理器会自动调用__enter__方法，然后进入运行时上下文环境，如果有as 从句，返回自身或另一个与运行时上下文相关的对象，值赋值给obj。
2、当with_body执行完毕退出with语句块或者with_body代码块出现异常，则会自动执行__exit__方法，并且会把对于的异常参数传递进来。
3、如果__exit__函数返回True。则with语句代码块不会显示的抛出异常，终止程序，如果返回None或者False，异常会被主动raise，并终止程序。

示例：自定义实现with open() 操作文件

# 使用上下文管理器实现上述文件操作
class MyOpenFile(object):
    def __init__(self, filename, mode='r'):
        self.file = open(filename, mode)

    def __enter__(self):
        return self.file

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.file.close()


with MyOpenFile('hello', 'w') as f:
    f.write("hello world")

1.3、上下文管理器的实际应用示例

需求：实现数据库资源操作包括创建资源和自动释放资源

import pymysql
class MyDataBase(object):
    def __init__(self, host, username, password, db, port):
        self.db = pymysql.connect(host=host, username=username, password=password, database=db, port=port)
        self.cursor = self.db.cursor()

    def query_one_data(self, sql):
        self.cursor.execute(sql)
        result = self.cursor.fetchone()
        return result

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cursor.close()
        self.db.close()

def test():
    sql = "select * from t_student where id = 1"
    with MyDataBase() as db:
        res = db.query_one_data(sql)
        print(res)

if __name__ == "__main__":
    test()

1.4、contextlib模块

contextlib 是 Python 的一个标准库模块，它提供了用于创建上下文管理器的装饰器和上下文管理器类。上下文管理器主要用于 with 语句，帮助自动管理资源，如文件操作、线程锁的获取与释放等，使得代码更加简洁、安全。

contextlib 模块中最常用的工具是 contextmanager 装饰器

1.4.1、contextmanager 装饰器

contextmanager 装饰器用于将一个简单的生成器函数转换成上下文管理器。这允许你避免创建一个类或单独的方法来实现 __enter__() 和 __exit__() 方法。
被装饰的函数被调用的时候必须返回一个生成器（被装饰器修饰的函数必须是一个生成器），而且这个生成器只生成一个值，如果有as的话，该值将绑定到with语句as子句的目标obj中。

示例1：

from contextlib import contextmanager  

@contextmanager  
def my_context():  
    print('Enter')  
    try:  
        yield  
    finally:  
        print('Exit')  

with my_context():  
    print('Inside the context')

根据输出结果，可以知道执行的流程：

1、先输出yield前的输出语句；
2、然后再是my_context()函数的输出语句，
3、最后是yield后面的输出语句。

等价于下面的写法：

def __enter__(self):
    print('Enter')


def __exit__(self, exc_type, exc_val, exc_tb):
    print('Exit')

示例2，使用contextmanager 装饰器实现1.3中数据库的管理操作：

import contextlib, pymysql
class MyDataBase(object):
    def __init__(self, host, username, password, db, port):
        self.db = pymysql.connect(host=host, username=username, password=password, database=db, port=port)
        self.cursor = self.db.cursor()

    def query_one_data(self, sql):
        self.cursor.execute(sql)
        result = self.cursor.fetchone()
        return result

@contextlib.contextmanager
def db_query_one():
    db = MyDataBase("127.0.0.1", "test", "123456", "student", 3306)
    yield db

def test():
    sql = "select * from t_student where id = 1"
    with db_query_one() as db:
       res =  db.query_one_data(sql)
       print(res)