Python with 语句的深入理解:优雅处理资源管理 @contextmanager

开发 前端
大家都用过 with open() as f 来读写文件,但可能较少去实现自己的 context manager。今天我们就通过几个实用场景,来深入理解这个优雅的语法特性。

大家都用过 with open() as f 来读写文件,但可能较少去实现自己的 context manager。今天我们就通过几个实用场景,来深入理解这个优雅的语法特性。

你一定用过:优雅处理资源管理

在 Python 中,如果不正确关闭文件句柄,可能带来严重后果:

# 错误示例
f = open('huge_file.txt')
content = f.read()
# 忘记调用 f.close()

# 潜在问题:
# 1. 文件句柄泄露:操作系统能打开的文件数是有限的
# 2. 数据丢失:写入的数据可能还在缓冲区,未真正写入磁盘
# 3. 文件锁定:其他程序可能无法访问该文件

这就是为什么我们推荐使用 with 语句:

with open('huge_file.txt') as f:
    content = f.read()
# 这里自动调用了 f.close(),即使发生异常也会关闭

那么,为什么使用了 with 可以自动调用 f.close() 呢?

从一个数据分析场景说起

假设你正在处理大量临时数据文件,下载后需要及时清理以节省磁盘空间:

def process_data():
    # 未使用 with 的写法
    try:
        data = download_large_file()
        result = analyze(data)
        cleanup_temp_files()
        return result
    except Exception as e:
        cleanup_temp_files()
        raise e

这种写法有几个问题:

  • cleanup 逻辑重复了
  • 如果中间加入 return ,容易忘记cleanup
  • 代码结构不够优雅

让我们改用 context manager 的方式:

class DataManager:
    def __enter__(self):
        self.data = download_large_file()
        return self.data

    def __exit__(self, exc_type, exc_value, traceback):
        cleanup_temp_files()
        return False  # 不吞掉异常

def process_data():
    with DataManager() as data:
        return analyze(data)  # 自动cleanup,更简洁

如上,当我们定义了 __enter__ 和 __exit__ 方法,Python 会在使用 with 语句时自动调用 __enter__,离开 with 语句时定义的作用域时自动调用 __exit__。

__exit__方法的返回值决定了异常是否会被"吞掉"(suppressed):

  1. 如果 __exit__ 返回 True :

如果在上下文管理器块中发生了异常,这个异常会被抑制

程序会继续正常执行,就像没有发生异常一样

  1. 如果 __exit__ 返回 False 或 None (默认):

异常会被重新抛出

程序会按照正常的异常处理流程执行

常见应用场景

1. 资源管理

with open('huge_file.txt') as f:
    content = f.read()

除了文件操作,还包括:

  • 数据库连接
  • 网络连接
  • 临时文件处理

2. 代码计时器

class Timer:
    def __enter__(self):
        self.start = time.time()  # 步骤1:进入 with 代码块时执行
        return self
        
    def __exit__(self, *args):
        self.end = time.time()    # 步骤3:离开 with 代码块时执行
        print(f'耗时: {self.end - self.start:.2f}秒')

# 使用示例
with Timer():
    time.sleep(1.5)  # 步骤2:执行 with 代码块内的代码
    # 步骤3会在这里自动执行,即使发生异常也会执行

3. 线程锁

from threading import Lock

class SafeCounter:
    def __init__(self):
        self._counter = 0
        self._lock = Lock()
    
    @property
    def counter(self):
        with self._lock:  # 自动加锁解锁
            return self._counter

@contextmanager 装饰器解析

除了定义类,还可以用装饰器 @contextmanager 来创建 context manager 。

当我们使用 @contextmanager 装饰一个生成器函数时,装饰器会:

  1. 创建一个新的类,实现 __enter__ 和 __exit__ 方法
  2. 将我们的生成器函数分成三部分:
  • yield 之前的代码放入 __enter__
  • yield 的值作为 __enter__ 的返回值
  • yield 之后的代码放入 __exit__

例如:

import os
from contextlib import contextmanager
import time

# 方式1:使用 @contextmanager 装饰器
@contextmanager
def temp_file(filename):
    # __enter__ 部分
    print(f"创建临时文件: {filename}")
    with open(filename, 'w') as f:
        f.write('一些临时数据')
    
    try:
        yield filename  # 返回值
    finally:
        # __exit__ 部分
        print(f"清理临时文件: {filename}")
        os.remove(filename)

# 方式2:使用传统的类实现
class TempFileManager:
    def __init__(self, filename):
        self.filename = filename
    
    def __enter__(self):
        print(f"创建临时文件: {self.filename}")
        with open(self.filename, 'w') as f:
            f.write('一些临时数据')
        return self.filename
    
    def __exit__(self, exc_type, exc_value, traceback):
        print(f"清理临时文件: {self.filename}")
        os.remove(self.filename)
        return False

# 测试代码
def process_file(filepath):
    print(f"处理文件: {filepath}")
    time.sleep(1)  # 模拟一些处理过程
    if "error" in filepath:
        raise ValueError("发现错误文件名!")

def test_context_manager():
    print("\n1. 测试 @contextmanager 装饰器版本:")
    try:
        with temp_file("test1.txt") as f:
            process_file(f)
        print("正常完成")
    except ValueError as e:
        print(f"捕获到异常: {e}")

    print("\n2. 测试类实现版本:")
    try:
        with TempFileManager("test2.txt") as f:
            process_file(f)
        print("正常完成")
    except ValueError as e:
        print(f"捕获到异常: {e}")

    print("\n3. 测试异常情况:")
    try:
        with temp_file("error.txt") as f:
            process_file(f)
        print("正常完成")
    except ValueError as e:
        print(f"捕获到异常: {e}")

if __name__ == "__main__":
    test_context_manager()

输出如下:

1. 测试 @contextmanager 装饰器版本:
创建临时文件: test1.txt
处理文件: test1.txt
清理临时文件: test1.txt
正常完成

2. 测试类实现版本:
创建临时文件: test2.txt
处理文件: test2.txt
清理临时文件: test2.txt
正常完成

3. 测试异常情况:
创建临时文件: error.txt
处理文件: error.txt
清理临时文件: error.txt
捕获到异常: 发现错误文件名!

高级用法:异常处理

__exit__ 方法可以优雅处理异常:

import sqlite3
import time
from contextlib import contextmanager

class Transaction:
    def __init__(self, db_path):
        self.db_path = db_path

    def __enter__(self):
        print("开始事务...")
        self.conn = sqlite3.connect(self.db_path)
        self.conn.execute('BEGIN TRANSACTION')
        return self.conn
        
    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            print("提交事务...")
            self.conn.commit()
        else:
            print(f"回滚事务... 异常: {exc_type.__name__}: {exc_value}")
            self.conn.rollback()
        self.conn.close()
        return False  # 不吞掉异常

# 为了对比,我们也实现一个装饰器版本
@contextmanager
def transaction(db_path):
    print("开始事务...")
    conn = sqlite3.connect(db_path)
    conn.execute('BEGIN TRANSACTION')
    try:
        yield conn
        print("提交事务...")
        conn.commit()
    except Exception as e:
        print(f"回滚事务... 异常: {type(e).__name__}: {e}")
        conn.rollback()
        raise  # 重新抛出异常
    finally:
        conn.close()

def init_db(db_path):
    """初始化数据库"""
    conn = sqlite3.connect(db_path)
    conn.execute('''
        CREATE TABLE IF NOT EXISTS accounts (
            id INTEGER PRIMARY KEY,
            name TEXT,
            balance REAL
        )
    ''')
    # 插入初始数据
    conn.execute('DELETE FROM accounts')  # 清空旧数据
    conn.execute('INSERT INTO accounts (name, balance) VALUES (?, ?)', ('Alice', 1000))
    conn.execute('INSERT INTO accounts (name, balance) VALUES (?, ?)', ('Bob', 1000))
    conn.commit()
    conn.close()

def transfer_money(conn, from_name, to_name, amount):
    """转账操作"""
    print(f"转账: {from_name} -> {to_name}, 金额: {amount}")
    
    # 模拟一些延迟,便于观察
    time.sleep(1)
    
    # 扣款
    cursor = conn.execute(
        'UPDATE accounts SET balance = balance - ? WHERE name = ? AND balance >= ?',
        (amount, from_name, amount)
    )
    if cursor.rowcount == 0:
        raise ValueError(f"{from_name} 余额不足或账户不存在!")
    
    # 模拟可能的错误情况
    if to_name == "ErrorUser":
        raise ValueError("目标账户不存在!")
    
    # 入账
    conn.execute(
        'UPDATE accounts SET balance = balance + ? WHERE name = ?',
        (amount, to_name)
    )

def show_balances(db_path):
    """显示所有账户余额"""
    conn = sqlite3.connect(db_path)
    cursor = conn.execute('SELECT name, balance FROM accounts')
    print("\n当前余额:")
    for name, balance in cursor:
        print(f"{name}: {balance}")
    conn.close()

def test_transactions():
    db_path = "test_transactions.db"
    init_db(db_path)
    
    print("\n1. 测试正常转账:")
    try:
        with Transaction(db_path) as conn:
            transfer_money(conn, "Alice", "Bob", 300)
        print("转账成功!")
    except Exception as e:
        print(f"转账失败: {e}")
    show_balances(db_path)

    print("\n2. 测试余额不足:")
    try:
        with Transaction(db_path) as conn:
            transfer_money(conn, "Alice", "Bob", 2000)
        print("转账成功!")
    except Exception as e:
        print(f"转账失败: {e}")
    show_balances(db_path)

    print("\n3. 测试无效账户:")
    try:
        with Transaction(db_path) as conn:
            transfer_money(conn, "Alice", "ErrorUser", 100)
        print("转账成功!")
    except Exception as e:
        print(f"转账失败: {e}")
    show_balances(db_path)

    print("\n4. 使用装饰器版本测试:")
    try:
        with transaction(db_path) as conn:
            transfer_money(conn, "Bob", "Alice", 200)
        print("转账成功!")
    except Exception as e:
        print(f"转账失败: {e}")
    show_balances(db_path)

if __name__ == "__main__":
    test_transactions()

输出如下:

1. 测试正常转账:
开始事务...
转账: Alice -> Bob, 金额: 300
提交事务...
转账成功!

当前余额:
Alice: 700.0
Bob: 1300.0

2. 测试余额不足:
开始事务...
转账: Alice -> Bob, 金额: 2000
回滚事务... 异常: ValueError: Alice 余额不足或账户不存在!
转账失败: Alice 余额不足或账户不存在!

当前余额:
Alice: 700.0
Bob: 1300.0

3. 测试无效账户:
开始事务...
转账: Alice -> ErrorUser, 金额: 100
回滚事务... 异常: ValueError: 目标账户不存在!
转账失败: 目标账户不存在!

当前余额:
Alice: 700.0
Bob: 1300.0

4. 使用装饰器版本测试:
开始事务...
转账: Bob -> Alice, 金额: 200
提交事务...
转账成功!

当前余额:
Alice: 900.0
Bob: 1100.0

实用建议

  1. 及时清理:__exit__ 确保资源释放
  2. 异常透明:通常返回 False,让异常继续传播
  3. 功能单一:一个 context manager 只做一件事
  4. 考虑可组合:多个 with 可以组合使用

小结

with 语句是 Python 中非常优雅的特性,善用它可以:

  • 自动管理资源
  • 简化异常处理
  • 提高代码可读性

建议大家在处理需要配对操作的场景(开启/关闭、加锁/解锁、创建/删除等)时,优先考虑使用 with 语句。

责任编辑:武晓燕 来源: Piper蛋窝
相关推荐

2024-07-18 10:12:04

2022-11-09 08:12:07

2023-12-04 08:46:40

Go标准库

2013-06-20 10:25:56

2024-10-15 15:58:11

2024-10-21 08:08:56

2016-12-08 15:36:59

HashMap数据结构hash函数

2010-06-01 15:25:27

JavaCLASSPATH

2020-07-21 08:26:08

SpringSecurity过滤器

2012-11-27 10:53:24

CloudPhysicvSphere

2012-03-27 22:22:51

iMC基础IT资源管理

2018-01-22 17:02:48

Python字符编码ASCII

2011-04-28 11:01:40

Android消息处理LooperHandler

2012-11-22 10:11:16

LispLisp教程

2018-11-30 10:00:53

Python字符串编程语言

2016-08-31 15:50:50

PythonThreadLocal变量

2016-11-07 21:59:52

threadpython

2016-08-31 15:41:19

PythonThreadLoca变量

2020-09-23 10:00:26

Redis数据库命令

2017-01-10 08:48:21

点赞
收藏

51CTO技术栈公众号