059、上下文管理器with 语句的原理、contextlib 装饰器与嵌套资源管理一个让我熬夜到凌晨三点的Bug去年接手一个数据管道项目每天凌晨定时从FTP拉取文件解压后写入数据库。上线一周一切正常。直到某天早上运维群里炸了——数据库连接池耗尽所有写入任务全部挂起。我翻了一上午日志发现一个诡异现象每次任务执行后数据库连接数都会增加几个但永远不会减少。更奇怪的是FTP连接也经常超时明明代码里写了ftp.quit()。排查到最后问题出在一个看似人畜无害的代码片段上defprocess_file():ftpFTP(data.company.com)ftp.login(user,pass)fileftp.retrbinary(RETR report.csv,open(report.csv,wb).write)# 处理文件...ftp.quit()returnresult你看出问题了吗如果retrbinary过程中抛出异常ftp.quit()永远不会执行。更致命的是数据库连接也是类似写法——try...finally嵌套了三层中间某个异常导致finally块被跳过。这就是为什么我们需要上下文管理器。不是因为它“优雅”而是因为它能救命。with语句到底干了什么很多人以为with open(...) as f只是语法糖省去了写f.close()的麻烦。这种理解太天真了。with语句的核心是协议——上下文管理协议。它由两个魔法方法组成__enter__和__exit__。classManagedFile:def__init__(self,filename,mode):self.filenamefilename self.modemodedef__enter__(self):# 这里踩过坑__enter__必须返回资源对象self.fileopen(self.filename,self.mode)returnself.file# 这个返回值会赋给as后面的变量def__exit__(self,exc_type,exc_val,exc_tb):# exc_type: 异常类型没有异常时为None# exc_val: 异常实例# exc_tb: 回溯信息ifself.file:self.file.close()# 别这样写return True会吞掉异常调试时你会疯掉# return True # 除非你明确知道要抑制异常returnFalse# 默认行为传播异常执行流程是这样的调用__enter__获取资源执行with块内的代码无论代码是否抛出异常都会调用__exit__如果__exit__返回True异常被抑制返回False异常继续传播这个机制保证了资源释放的确定性——这正是我那个FTP连接问题的解药。用contextlib装饰器偷懒每次写__enter__和__exit__太啰嗦了。Python标准库提供了contextlib里面有几个实用工具。contextmanager装饰器这是我最常用的方式把生成器函数变成上下文管理器fromcontextlibimportcontextmanagercontextmanagerdefmanaged_ftp(host,user,password):# yield之前的代码相当于__enter__ftpFTP(host)ftp.login(user,password)try:yieldftp# 这个yield的值会赋给as后面的变量finally:# yield之后的代码相当于__exit__# 注意这里用finally保证即使yield内部异常也会执行ftp.quit()使用起来很直观withmanaged_ftp(data.company.com,user,pass)asftp:ftp.retrbinary(RETR report.csv,open(report.csv,wb).write)# 即使这里抛出异常ftp.quit()也会执行踩坑提醒contextmanager装饰的函数必须是一个生成器有yield而且yield只能出现一次。如果你在yield后面写了多个yieldPython会抛出RuntimeError: generator didnt stop。closing()——处理没有上下文管理的对象有些老旧的库对象只有close()方法没有实现上下文协议。closing()可以包装它们fromcontextlibimportclosingimporturllib.requestwithclosing(urllib.request.urlopen(http://example.com))asresponse:dataresponse.read()# 退出时自动调用response.close()suppress()——优雅地忽略特定异常别再用try: ... except: pass了那会连KeyboardInterrupt都吞掉fromcontextlibimportsuppressimportos# 删除文件如果不存在就忽略withsuppress(FileNotFoundError):os.remove(temp.txt)# 等价于# try:# os.remove(temp.txt)# except FileNotFoundError:# passredirect_stdout/redirect_stderr——临时重定向输出调试时很有用特别是处理第三方库的打印日志fromcontextlibimportredirect_stdoutimportio fio.StringIO()withredirect_stdout(f):print(这段文字不会打印到控制台)help(print)# help()的输出也会被重定向outputf.getvalue()嵌套资源管理——真正的考验回到我那个数据库连接池耗尽的问题。当时代码结构大概是这样的# 别这样写——嵌套try...finally地狱ftpNoneconnNonetry:ftpFTP(host)ftp.login(user,pass)try:connget_db_connection()cursorconn.cursor()try:dataftp.retrlines(LIST)cursor.execute(INSERT INTO ...,data)finally:cursor.close()finally:conn.close()finally:ifftp:ftp.quit()这种写法的问题如果cursor.execute()抛出异常cursor.close()会执行但conn.close()和ftp.quit()也会执行。看起来没问题但如果在conn.close()中又抛出了异常ftp.quit()就被跳过了。正确的做法是用多个with语句# 推荐写法——多个with可以写在一行withFTP(host)asftp,get_db_connection()asconn:ftp.login(user,pass)withconn.cursor()ascursor:dataftp.retrlines(LIST)cursor.execute(INSERT INTO ...,data)Python 3.1支持多个上下文管理器用逗号分隔。Python 3.10还支持括号换行with(FTP(host)asftp,get_db_connection()asconn,conn.cursor()ascursor):ftp.login(user,pass)dataftp.retrlines(LIST)cursor.execute(INSERT INTO ...,data)每个with语句都会独立管理自己的资源一个异常不会影响其他资源的释放。一个生产级的上下文管理器模板这是我项目中实际使用的数据库连接管理器处理了连接池、超时、重试等场景fromcontextlibimportcontextmanagerfromtypingimportGenerator,Optionalimporttimefrommy_db_libimportConnectionPool,DatabaseErrorclassDatabaseSession:数据库会话管理器支持自动重连和超时控制def__init__(self,pool:ConnectionPool,timeout:int30):self.poolpool self.timeouttimeout self.conn:Optional[Connection]Noneself.start_time:Optional[float]Nonedef__enter__(self):self.start_timetime.time()# 这里踩过坑连接池可能返回Noneself.connself.pool.get_connection(timeoutself.timeout)ifself.connisNone:raiseConnectionError(无法从连接池获取连接)returnself.conndef__exit__(self,exc_type,exc_val,exc_tb):ifself.connisNone:returnFalseelapsedtime.time()-self.start_timeifexc_typeisnotNone:# 发生异常时回滚事务try:self.conn.rollback()exceptException:pass# 回滚失败也不能影响资源释放# 如果执行时间过长可能是死锁关闭连接而不是归还ifelapsedself.timeout:self.conn.close()self.connNonereturnFalse# 正常情况归还连接池try:self.pool.return_connection(self.conn)exceptException:self.conn.close()finally:self.connNonereturnFalse# 不抑制异常个人经验总结永远不要手动管理资源。try...finally嵌套超过两层代码就不可维护了。用with语句一个资源一个with。contextmanager是你的好朋友。但记住yield前面的代码是__enter__yield后面是__exit__中间用try...finally包裹yield确保异常时也能执行清理。别滥用suppress()。只忽略你确定无害的异常比如FileNotFoundError。永远不要用suppress(Exception)那等于关掉了所有错误报告。自定义上下文管理器时__exit__的返回值要谨慎。返回True会吞掉异常除非你明确知道要这么做比如实现了重试逻辑否则返回False。测试你的上下文管理器。写一个单元测试在with块内故意抛出异常验证资源是否被正确释放。这个测试能救你于水火。Python 3.10的括号语法让嵌套管理变得清晰推荐使用。但注意括号内的每个with项都是独立的上下文管理器它们的__enter__按从左到右顺序调用__exit__按从右到左顺序调用——这符合直觉因为最内层的资源应该最先释放。那个让我熬夜的Bug最后修复方案就是把所有资源管理改成了with语句。上线后连接池再也没出过问题。有时候最优雅的解决方案就是最基础的方案——把资源管理交给语言本身而不是靠自己的记忆力。