Python魔法模块之contextlib

引言

我们在操作文件时最常用的就是使用with上下文管理器，这样会让代码的可读性更强而且错误更少，例如：

with open('/tmp/a.txt', a) as file_obj:
    file_obj.write("hello carson")

按照上述这样写的好处在于，在执行完毕缩进代码块后会自动关闭文件。同样的例子还有threading.Lock，如果不使用with，需要这样写：

import threading
lock = threading.Lock()

lock.acquire()
try:
    my_list.append(item)
finally:
    lock.release()

如果使用with，那就会非常简单：

with lock:
    my_list.append(item)

创建上下文管理(类实现)

创建上下文管理实际就是创建一个类，添加__enter__和__exit__方法。下面我们来实现open的上下文管理功能：

class OpenContext(object):

    def __init__(self, filename, mode):
        self.fp = open(filename, mode)

    def __enter__(self):
        return self.fp

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.fp.close()

        
with OpenContext('/tmp/a.txt', 'a') as file_obj:
    file_obj.write("hello 6666")

使用@contextmanager创建

上面我们自定义上下文管理器确实很方便，但是Python标准库还提供了更加易用的上下文管理器工具模块contextlib，它是通过生成器实现的，我们不需要再创建类以及__enter__和__exit__这两个特俗的方法：

from contextlib import contextmanager

@contextmanager
def make_open_context(filename, mode):
    fp = open(filename, mode)
    try:
        yield fp
    finally:
        fp.close()

with make_open_context('/tmp/a.txt', 'a') as file_obj:
    file_obj.write("hello carson666")

在上文中，yield关键词把上下文分割成两部分：yield之前就是__init__中的代码块；yield之后其实就是__exit__中的代码块，yield生成的值会绑定到with语句as子句中的变量，例如在上面的例子中，yield生成的值是文件句柄对象fp，在下面的with语句中，会将fp和file_obj绑定到一起，也就是说file_obj此时就是一个文件句柄对象，那么它就可以操作文件了，因此就可以调用file_obj.write("hello carson666")，另外要注意的是如果yield没有生成值，那么在with语句中就不需要写as子句了，后面会结合具体的例子详解。

案例一：

# _*_ coding:utf-8 _*_
from contextlib import contextmanager

"""
contextmanager给了我们一个机会，即将原来不是上下文管理器的类变成了一个
上下文管理器，例如这里的MyResource类
"""
class MyResource:
    def query(self):
        print("query data")

@contextmanager
def make_myresource():
    print("connect to resource")
    yield MyResource()
    print("connect to resource")

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

上面的例子就充分体现了contextmanager的强大作用，将一个不是上下问管理器的类 MyResource变成了一个上下文管理器，这样做的好处在于，我们就可以在执行真正的核心代码之前可以执行一部分代码，然后在执行完毕后，又可以执行一部分代码，这种场景在实际需求中还是很常见的。上面yield MyResource() 生成了一个实例对象，然后我们可以在with语句中调用类中的方法。看看最终的打印结果：

上面这样写的好处还有：假如MyResource是Flask或者其他第三方插件提供给我们的类库，如果使用自定义上下文管理器，那么就要使用手动去修改源码，在原代码的基础上添加enter和exit方法，这样做肯定不合适；现在我们可以在类MyResource的外部使用contextmanager将该类包装成为一个上下文管理器，这样既可以调用类中的方法，又可以在执行核心代码前后再执行一些相关的语句。

案例二

我现在想打印 《且将生活一饮而尽》 也就是说我们要打印《 和 》；这就体现了上下文管理器的一个用法：我仅仅就是需要在我的核心代码前面和后面各执行一段代码而已：

# _*_ coding:utf-8 _*_

from contextlib import contextmanager

@contextmanager
def book_mark():
    print('《', end="")
    yield
    print('》', end="")

with book_mark():
    # 核心代码
    print('且将生活一饮而尽', end="")

案例三：

看看实际的例子，我们通常在SQLAlchemy中使用db.session.commit()，既然有commit，那就需要做一个事务的处理。因此我们需要使用try except来处理异常，如下图所示：

一般在应用程序中，我们都有很多个db.session.commit(),那如果都要使用try来处理异常，那就太麻烦了。我们需要做的是在核心代码之前使用try,然后在核心代码执行完毕之后，加上except。因此这就使用到了上下文管理器，此时这个db是一个第三方类库SQLAlchemy，那我们如何去为它新增加一个方法呢？那我们就继承SQLAlchemy，首先导入，然后取名字

上面我们就为SQLAlchemy新增了一个auto_commit方法，主要实现的是自动commit()和rollback()；并将其变成了一个上下文管理器。

那么原始的代码就可以变成如下的：

如下，我们在保存user的时候也使用到了db.session.commit()；所以也可以改写