python 上下文管理器contextlib.ContextManager

1 模块简介

在数年前，Python 2.5 加入了一个非常特殊的关键字，就是with。with语句允许开发者创建上下文管理器。什么是上下文管理器？上下文管理器就是允许你可以自动地开始和结束一些事情。例如，你可能想要打开一个文件，然后写入一些内容，最后再关闭文件。这或许就是上下文管理器中一个最经典的示例。事实上，当你利用with语句打开一个文件时，Python替你自动创建了一个上下文管理器。

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

如果你使用的是Python 2.4，你不得不以一种老的方式来完成这个任务

f_obj = open("test/test.txt","w")
f_obj.write("hello")
f_obj.close()

下文管理器背后工作的机制是使用Python的方法：__enter__和__exit__。让我们尝试着去创建我们的上下文管理器，以此来了解上下文管理器是如何工作的。

2 模块使用

2.0 一个简单的demo

import contextlib
import time

@contextlib.contextmanager
def timeit(title):
    print('1...')
    start = time.time()
    yield
    print('2...')
    end = time.time()
    usedTime = (end - start) * 1000
    print('Use time %d ms' % usedTime)


with timeit(1):
    print('3...')
    time.sleep(1)

with timeit(2):
    print('4...')
    time.sleep(2)

输出结果：

1...
3...
2...
Use time 1001 ms
1...
4...
2...
Use time 2002 ms

2.1 创建一个上下文管理器类

与其继续使用Python打开文件这个例子，不如我们创建一个上下文管理器，这个上下文管理器将会创建一个SQLite数据库连接，当任务处理完毕，将会将其关闭。下面就是一个简单的示例。

import sqlite3

class DataConn:
    def __init__(self,db_name):
        self.db_name = db_name

    def __enter__(self):
        self.conn = sqlite3.connect(self.db_name)
        return self.conn

    def __exit__(self,exc_type,exc_val,exc_tb):
        self.conn.close()
        if exc_val:
            raise

if __name__ == "__main__":
    # with 装饰器的底层原理实际上是__enter__和__exit__  实现的。
    db = "test/test.db"
    with DataConn(db) as conn:
        cursor = conn.cursor()

2.2 利用contextlib创建一个上下文管理器

Python 2.5 不仅仅添加了with语句，它也添加了contextlib模块。这就允许我们使用contextlib的contextmanager函数作为装饰器，来创建一个上下文管理器。让我们尝试着用它来创建一个上下文管理器，用于打开和关闭文件

from contextlib import contextmanager

@contextmanager
def file_open(path):
    try:
        f_obj = open(path,"w")
        yield f_obj
    except OSError:
        print("We had an error!")
    finally:
        print("Closing file")
        f_obj.close()

if __name__ == "__main__":
    with file_open("test/test.txt") as fobj:
        fobj.write("Testing context managers")

原文链接