[python] 上下文管理器

什么是上下文管理器

简单来说，上下文管理器的目的就是规定对象的使用范围，如果超出范围就采取相应“处理”；比如：

f = open('filename.txt')
data = f.read()
f.close()

打开一个文件，读取文件内容，关闭文件。正常的业务中，这样写是不够的，因为在操作资源的时候有可能出现错误，为了加强代码健壮性，需要经常改为：​​​​​​​

try:
    f = open('filename.txt')
    data = f.read()
except:
    pass
finally:
    f.close()

不管出现什么错误，最后都要关闭和释放资源f.close()，为了代码更具可读性，且不容易出错，就会用with实现上下文管理器：​​​​​​​

with open('filename.txt') as f:
    data = f.read()

先理清几个概念：​​​​​​​

1. 上下文表达式：with open('filename.txt') as f:2. 上下文管理器：open('filename.txt')3. f 不是上下文管理器，应该是资源对象，是上下文管理器返回的对象

实现上下文管理器

要自己实现这样一个上下文管理，要先知道上下文管理协议。简单点说，就是在一个类里，实现了__enter__和__exit__的方法，这个类的实例就是一个上下文管理器。​​​​​​​

class MyResource:
    def __enter__(self):
        print('opening resource......')
        return self

    def operate(self):
        print('doing something......')

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('closing resource......')


with MyResource() as r:
    r.operate()

当打开资源的时候进入__enter__方法，return的对象会赋给as对象，即这里return回self对象会赋值给r；不管有没有出错，执行完with代码块后都会进入到__exit__方法，如果出错，__exit__的参数就会被赋值。

业务实战

实现一个连接mysql数据库的上下文管理器，就以上面的代码为基础​​​​​​​

import pymysql


class MyResource:
    
    def __init__(self, database):
        self.database = database
    
    def __enter__(self):
        self.conn = pymysql.connect(
            host='localhost',
            port=3306,
            user='root',
            password='xxxxxx',
            database=self.database,
            charset='utf8'
        )
        self.cursor = self.conn.cursor()
        return self.cursor

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


with MyResource('datatbase') as db:
    db.execute('update test set name="rhys" where id=1')

用with实例化MyResource对象，进入__enter__函数连接数据库，返回游标cursor给db，db.execute进行update操作，在退出with代码块的时候进入__exit__函数，进行commit操作再关闭游标和连接。

contextmanager实现上下文管理器

还有另一种方式实现一个上下问管理器，就是用contextlib里的contextmanager装饰器：​​​​​​​

from contextlib import contextmanager
import pymysql


@contextmanager
def mysqldb(database):
    try:
        conn = pymysql.connect(
            host='localhost',
            port=3306,
            user='root',
            password='xxxxxx',
            database=database,
            charset='utf8'
        )
        cursor = conn.cursor()
        yield cursor
        conn.commit()
    except Exception as e:
        print(e)
    finally:
        cursor.close()
        conn.close()


with mysqldb('database') as db:
    db.execute('update test set name="rhys" where id=1')

被contextmanager装饰的函数会成为一个上下文管理器，用yield返回cursor给as后面的db，执行完update操作退出with代码块的时候再回到yield的位置，执行之后的代码。

个人更喜欢使用contexmanager来实现上下问管理器，代码更简便更容易编写，也更好理解。

 

关注公众号：日常bug，适合技术点滴积累，利用琐碎时间学习技术的人。