Python 自定义上下文管理器 和 with 的使用

with语句

Python提供了 with 语句的这种写法，既简单又安全，并且 with 语句执行完成以后自动调用关闭文件操作，即使出现异常也会自动调用关闭文件操作。

# 只读方式打开文件，写入”sss“报错，但 文件仍然会关闭
with open("1.txt", mode="r") as file:
    file.write("sss")

上下文管理器--类

一个类只要实现了__enter__() 和__exit__() 这个两个方法，通过该类创建的对象我们就称之为上下文管理器。

定义一个File类，实现 __enter__() 和 __exit__()方法，然后使用 with 语句来完成操作文件， 示例代码:

# coding=utf8
"""在类里面，实现 __enter__() 和 __exit__()方法, 就是上下文管理器"""

"""使用上下文管理器 + with 进行文件操作"""


# 自定义上下文管理器， 创建类
class File:
    def __init__(self, file_name, file_mode):
        self.file_name = file_name
        self.file_mode = file_mode

    # 上文方法：负责返回操作对象资源，比如：文件对象，数据库连接对象（都有关闭动作）
    def __enter__(self):
        # self.file    将变量变成对象属性
        self.file = open(self.file_name, self.file_mode)
        return self.file

    # 下文方法：负责释放对象资源，比如：文件对象，数据库连接对象（都有关闭动作）
    # 当with执行完后，自动执行__exit__方法  比如 关闭文件
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("over")
        self.file.close()


# with 语句结合上下文管理器使用
# 类（方法）， 会执行init 方法， 传入file_name 和 file_mode 两参数
# as 接上文管理器的返回对象  aaa = self.file
with File("1.txt", "r") as aaa:
    file_data = aaa.read()
    # 只读模式打开，就算写文件报错，也要执行__exit__(), 关闭文件
    # aaa.write("sss")
print(file_data)

定义一个Mysql类，实现 __enter__() 和 __exit__()方法，然后使用 with 语句来完成操作数据库存， 示例代码:

# coding=utf8
"""使用上下文管理器 + with 进行 Mysql 操作"""
import pymysql


# 定义上下文管理器，创建类
class Mysql:
    def __init__(self, host, port, user, password, database, charset):
        self.sql_host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.charset = charset

    # 上文方法：负责返回操作对象资源
    def __enter__(self):
        # self.db 将变量变成对象属性, 连接数据库
        self.db = pymysql.connect(host=self.sql_host, port=self.port, user=self.user,
                                  password=self.password, database=self.database,
                                  charset=self.charset)
        # self.cursor 将变量变成对象属性, 游标
        self.cursor = self.db.cursor()
        return self.cursor

    # 下文方法：负责释放对象资源
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cursor.close()
        self.db.close()
        print("Mysql访问完成~")


# Mysql() 方法，会执行 Mysql类的__init__属性，传入连接数据库的参数
# as 接上文管理器的返回对象  s 就是 上文 方法中的 self.cursor ,  游标
with Mysql("127.0.0.1", 3306, "root", "123456",
           "fujifilm", "") as s:
    sql = "select * from sa_list;"
    s.execute(sql)      # s = self.cursor
    result = s.fetchall()       # s = self.cursor
    print(result)

上下文管理器-函数-@contextmanager

假如想要让一个函数成为上下文管理器，Python 还提供了一个 @contextmanager 的装饰器，更进一步简化了上下文管理器的实现方式。通过 yield 将函数分割成两部分，yield 上面的语句在 __enter__ 方法中执行，yield 下面的语句在 __exit__ 方法中执行，紧跟在 yield 后面的参数是函数的返回值。

# coding=utf8
from contextlib import contextmanager
import pymysql

"""使用上下文管理器 + with 进行 文件 操作"""


# 加上@contextmanager 装饰器代码，下面函数创建的对象就是一个上下文管理器
@contextmanager
def my_open(file_name, file_mode):
    try:
        file = open(file_name, file_mode)
        # yield 之前的代码可以认为是上文方法，负责返回操作对象资源
        yield file
    except Exception as e:
        print(e)
    finally:
        print("over")
        # yield 之后的代码可以认为是下文方法，负责释放操作对象资源
        file.close()


# 普通函数不能结合with语句使用，需要结合上下文管理器使用
with my_open("1.txt", file_mode="r") as f:
    # f 就是上文管理器中返回的 file
    data = f.read()
    print(data)


"""使用上下文管理器 + with 进行 Mysql 操作"""


# 加上@contextmanager 装饰器代码，下面函数创建的对象就是一个上下文管理器
@contextmanager
def Mysql_conn(host, port, user, password, database, charset):
    try:
        con = pymysql.connect(host=host, port=port, user=user,
                              password=password, database=database,
                              charset=charset)
        cursor = con.cursor()  # 游标
        yield cursor  # 用 yield 返回游标
    except Exception as e:
        print(e)
    finally:
        cursor.close()
        con.close()
        print("MySQL连接完成")


# 结合with语句
with Mysql_conn("localhost", 3306, "root", "123456", "fujifilm", "") as s:
    # s 就是函数上文方法中返回的 cursor 游标
    sql = "select * from sa_list;"
    s.execute(sql)
    result = s.fetchall()
    print(result)