承上启下继往开来,Python3上下文管理器(ContextManagers)与With关键字的迷思

by Liu Yue/2022-07-19

    在开发过程中,我们会经常面临的一个常见问题是如何正确管理外部资源,比如数据库、锁或者网络连接。稍不留意,程序将永久保留这些资源,即使我们不再需要它们。此类问题被称之为内存泄漏,因为每次在不关闭现有资源的情况下创建和打开给定资源的新实例时,可用内存都会减少。

    正确管理资源往往是一个棘手的问题,因为资源的使用往往需要进行善后工作。善后工作要求执行一些清理操作,例如关闭数据库、释放锁或关闭网络连接。如果忘记执行这些清理操作,就可能会浪费宝贵的系统资源,例如内存和网络带宽。

    背景

    譬如,当开发人员使用数据库时,可能会出现一个常见问题是程序不断创建新连接而不释放或重用它们。在这种情况下,数据库后端可以停止接受新连接。这可能需要管理员登录并手动终止这些陈旧的连接,以使数据库再次可用。

    以著名的ORM工具Peewee为例子:

pip3 install pymysql
pip3 install peewee

    当我们声明数据库实例之后,试图链接数据库:

from peewee import MySQLDatabase

db = MySQLDatabase('mytest', user='root', password='root',host='localhost', port=3306)

print(db.connect())

    程序输出:

True

    但如果重复的创建链接:

from peewee import MySQLDatabase

db = MySQLDatabase('mytest', user='root', password='root',host='localhost', port=3306)

print(db.connect())
print(db.connect())

    就会抛出异常:

Traceback (most recent call last):
File "/Users/liuyue/Downloads/upload/test/test.py", line 23, in <module>
print(db.connect())
File "/opt/homebrew/lib/python3.9/site-packages/peewee.py", line 3129, in connect
raise OperationalError('Connection already opened.')
peewee.OperationalError: Connection already opened.

    所以,需要手动关闭数据库链接:

from peewee import MySQLDatabase

db = MySQLDatabase('mytest', user='root', password='root',host='localhost', port=3306)

print(db.connect())
print(db.close())
print(db.connect())

    返回:

True
True
True

    但这样操作有一个潜在的问题,如果在调用connect的过程中,出现了异常进而导致后续代码无法继续执行,close方法无法被正常调用,因此数据库资源就会一直被该程序占用而无法被释放。

    继续改进:

from peewee import MySQLDatabase

db = MySQLDatabase('mytest', user='root', password='root',host='localhost', port=3306)

try:
print(db.connect())
except OperationalError:
print("Connection already opened.")
finally:
print(db.close())

    改进后的逻辑是对可能发生异常的代码处进行OperationalError异常捕获,使用 try/finally 语句,该语句表示如果在 try 代码块中程序出现了异常,后续代码就不再执行,而直接跳转到 except 代码块。而最终,finally块逻辑的代码被执行。因此,只要把 close方法放在 finally 代码块中,数据库链接就会被关闭。

    事实上,Peewee为我们提供了一种更加简洁、优雅的方式来操作数据库链接:

from peewee import MySQLDatabase

db = MySQLDatabase('mytest', user='root', password='root',host='localhost', port=3306)

with db.connection_context():
print("db is open")
print(db.is_closed())

    也就是使用with 关键字来进行操作,这里使用with开启数据库的上下文管理器,当程序离开with关键字的作用域时,系统会自动调用close方法,最终效果和上文的捕获OperationalError异常一致,系统会自动关闭数据库链接。

    上下文管理器(ContextManagers)

    那么Peewee底层是如何实现对数据库的自动关闭呢?那就是使用Python3内置的上下文管理器,在Python中,任何实现了 __enter__() 和 __exit__() 方法的对象都可称之为上下文管理器,上下文管理器对象可以使用 with 关键字:

from peewee import MySQLDatabase

db = MySQLDatabase('mytest', user='root', password='root',host='localhost', port=3306)


class Db:

def __init__(self):

self.db = MySQLDatabase('mytest', user='root', password='root',host='localhost', port=3306)

def __enter__(self):

self.db.connect()

def __exit__(self,*args):

self.db.close()

    __enter__() 方法负责打开数据库链接,__exit__() 方法负责处理一些善后工作,也就是关闭数据库链接。

    藉此,我们就可以使用with关键字对Db这个类对象进行调用了:

with Db() as db:
print("db is opening")

print(Db().db.is_closed())

    程序返回:

db is opening
True

    如此,我们就无需显性地调用 close方法了,由系统自动去调用,哪怕中间抛出异常 close方法理论上也会被调用。

    上下文语法糖

    Python3 还提供了一个基于上下文管理器的装饰器,更进一步简化了上下文管理器的实现方式。通过 生成器yield关键字将方法分割成两部分,yield 之前的语句在 __enter__ 方法中执行,yield 之后的语句在 __exit__ 方法中执行。紧跟在 yield 后面的值是函数的返回值:

from peewee import MySQLDatabase
from contextlib import contextmanager


@contextmanager
def mydb():
db = MySQLDatabase('mytest', user='root', password='root',host='localhost', port=3306)
yield db
db.close()

    随后通过with关键字调用contextmanager装饰后的方法:

with mydb() as db:
print("db is opening")

    与此同时,Peewee也贴心地帮我们将此装饰器封装了起来:

from peewee import MySQLDatabase

db = MySQLDatabase('mytest', user='root', password='root',host='localhost', port=3306)


@db.connection_context()
def mydb():
print("db is open")

mydb()

    看起来还不错。

    迷思:上下文管理一定可以善后吗?

    请别太笃定,是的,上下文管理器美则美矣,但却未尽善焉,在一些极端情况下,未必能够善后:

from peewee import MySQLDatabase

class Db:

def __init__(self):

self.db = MySQLDatabase('mytest', user='root', password='root',host='localhost', port=3306)

def __enter__(self):
print("connect")
self.db.connect()
exit(-1)

def __exit__(self,*args):
print("close")
self.db.close()

with Db() as db:
print("db is opening")

    程序返回:

connect

    当我们通过with关键字调用上下文管理器时,在__enter__方法内通过exit()方法强行关闭程序,过程中程序会立刻结束,并未进入到__exit__方法中执行关闭流程,也就是说,这个数据库链接并未被正确关闭。

    同理,当我们书写了finally关键字,理所当然的,finally代码块理论上一定会执行,但其实,也仅仅是理论上:

def gen(text):
try:
for line in text:
try:
yield int(line)
except:

pass
finally:
print('善后工作')

text = ['1', '', '2', '', '3']

if any(n > 1 for n in gen(text)):
print('Found a number')

print('并未善后')

    程序返回:

Exception ignored in: <generator object gen at 0x100e177b0>
Traceback (most recent call last):
File "/Users/liuyue/Downloads/upload/test/test.py", line 71, in <module>
if any(n > 1 for n in gen(text)):
RuntimeError: generator ignored GeneratorExit
Found a number
并未善后

    显而易见,当程序进入finally代码块之前,就立刻触发了一个生成器generator异常,当理论上要被捕获异常时程序被yield返回了原始状态,于是立刻退出,放弃了执行finally逻辑。

    所以,逻辑上,我们并不能指望上下文管理器每一次都能够帮我们“善后”,至少,在事情尚未收束的情况下,能够随机应变:

from peewee import MySQLDatabase

class Db:

def __init__(self):

self.db = MySQLDatabase('mytest', user='root', password='root',host='localhost', port=3306)

def __enter__(self):

if self.db.is_closeed():
print("connect")
self.db.connect()

def __exit__(self,*args):
print("close")
self.db.close()

with Db() as db:
print("db is opening")

print(Db().db.is_closed())

    结语

    使用With关键字操作上下文管理器可以更快捷地管理外部资源,同时能提高代码的健壮性和可读性,但在极端情况下,上下文管理器也并非万能,还是需要诸如轮询服务等托底保障方案。

    我的博客即将同步至腾讯云开发者社区,邀请大家一同入驻:https://cloud.tencent.com/developer/support-plan?invite_code=131vd9e43v16l