[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，適合技術點滴積累，利用瑣碎時間學習技術的人。