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")

原文鏈接