contextmanager: 上下文管理器类和上下文管理器装饰器

原文

一. 什么是上下文管理器

二. 自定义一个上下文管理器类:

三. 使用contextmanager

四. 一个例子, sqlalchemy: 数据库的自动提交和回滚

---

一. 什么是上下文管理器

上下文管理器是在Python2.5之后加入的功能，可以在方便的需要的时候比较精确地分配和释放资源, with便是上下文管理器的最广泛的应用, 比如:

with open("test/test.txt","w") as f:
  f.write("hello")

这上会比使用try:...finally:f.close方便的多.

二. 自定义一个上下文管理器类:

class MyResource:
    # __enter__ 返回的对象会被with语句中as后的变量接受
    def __enter__(self):
        print('connect to resource')
        return self

    def __exit__(self, exc_type, exc_value, tb):
        print('close resource conection')

    def query(self):
        print('query data')

类中有两个特殊的魔术方法:
__enter__: with语句中的代码块执行前, 会执行__enter__, 返回的值将赋值给with句中as后的变量.
__exit__: with语句中的代码块执行结束或出错, 会执行_exit__

比如以下代码:

with Myresource() as r:
    r.query()

的打印结果为:

connect to resource
query data
close resource conection

那么有没有一个简化定义的方法呢, python提供了一个装饰器contextmanager

三. 使用contextmanager

from contextlib import contextmanager
class MyResource:
    def query(self):
        print('query data')

@contextmanager
def make_myresource():
    print('start to connect')
    yield MyResource()
    print('end connect')
    pass

被装饰器装饰的函数分为三部分:
with语句中的代码块执行前执行函数中yield之前代码
yield返回的内容复制给as之后的变量
with代码块执行完毕后执行函数中yield之后的代码

比如下方代码:

with make_myresource() as r:
     r.query()

的结果为:

start to connect
query data
end connect

四. 一个例子, sqlalchemy: 数据库的自动提交和回滚

在编程中如果频繁的修改数据库, 一味的使用类似try:... except..: rollback() raise e其实是不太好的.

比如某一段的代码的是这样的:

   try:
        gift = Gift()
        gift.isbn = isbn
        ... 
        db.session.add(gift)
        db.session.commit()
    except Exception as e:
        db.session.rollback()
        raise e

为了达到使用with语句的目的, 我们可以重写db所属的类:

from flask_sqlalchemy import SQLAlchemy as _SQLALchemy
class SQLAlchemy(_SQLALchemy):
    @contextmanager
    def auto_commit(self):
        try:
            yield
            self.session.commit()
        except Exception as e:
            db.session.rollback()
            raise e

这时候, 在执行数据的修改的时候便可以:

 with db.auto_commit():
        gift = Gift()
        gift.isbn = isbndb.session.add(gift)
        db.session.add(gift)

with db.auto_commit():
    user = User()
    user.set_attrs(form.data)
    db.session.add(user)