python contextlib 上下文管理器

 

1、with操作符

在python中讀寫文件，可能需要這樣的代碼

try-finally讀寫文件

file_text = None
try:
    file_text = open('./text', 'r')
    print file_text.read()
except IOError, ex:
    traceback.print_exc()
finally:
    if file_text:
        file_text.close()

同樣，在python中使用線程鎖，可能需要這樣的代碼

try-finally線程鎖

lock = threading.Lock()
lock.acquire()
try:
    pass
except Exception, ex:
    traceback.print_exc()
finally:
    lock.release()

可能你會覺得這種寫法很不方便，python提供了with操作符，你可以這樣操作

with讀寫文件

with open('./text', 'r') as file_text:
    print file_text.read()

with線程鎖

with lock:
    pass

是不是方便多了。

其實，不只是lock和file可以使用with操作符。

實際上，任何對象，只要正確實現上下文管理，就可以使用with語句。實現上下文管理是通過 __enter__ 和 __exit__ 這兩個方法實現的。

 

2、上下文管理

上下文管理可以為我們屏蔽上下文的復雜性。例如，我們實現一個類Cat，實現其__enter__和__exit__方法。

__enter__(self)： 進入上下文管理器時調用此方法，其返回值將被放入with-as語句中as說明符指定的變量中。

__exit__(self,type,value,tb):離開上下文管理器調用此方法。如果有異常出現，type、value、tb分別為異常的類型、值和追蹤信息。如果沒有異常，

3個參數均設為None。此方法返回值為True或者False，分別指示被引發的異常得到了還是沒有得到處理。如果返回False，引發的異常會被傳遞出上下文。

如下。

class Cat(object):

    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print 'enter from Cat named %s' % self.name
        return self

    def hello(self):
        print 'hello, %s' % self.name

    def __exit__(self, exc_type, exc_val, exc_tb):
        print 'exit from Cat named %s' % self.name

執行，並打印結果

with Cat('Tom') as tom:
    tom.hello()

enter from Cat named Tom
hello, Tom
exit from Cat named Tom

這里我們執行as tom獲得了Cat類的一個實例，這是通過__enter__方法的返回得到的。

當然，我們可以管理多個，請注意進入和退出的順序。

with Cat('Tom') as tom, Cat('Harry') as harry:
    tom.hello()
    harry.hello()

enter from Cat named Tom
enter from Cat named Harry
hello, Tom
hello, Harry
exit from Cat named Harry
exit from Cat named Tom

 

3、contextmanager

可能你還是覺得實現__enter__和__exit__很麻煩。python提供了contextlib.contextmanager

讓我們重寫上面的例子，使用contextmanager

from contextlib import contextmanager as _contextmanager


@_contextmanager
def cat(name):
    print 'enter cat named %s' % name
    yield name
    print 'exit cat named %s' % name

執行，並打印結果

with cat('Kitty') as kitty:
    print 'hello, %s' % kitty

enter cat named Kitty
hello, Kitty
exit cat named Kitty

as后面的實例，是通過yield語句返回的。這里是返回了一個字符串。

當然，同樣支持管理多個實例

with cat('Kitty') as kitty, cat('Tom') as tom:
    print 'hello, %s' % kitty
    print 'hello, %s' % tom

enter cat named Kitty
enter cat named Tom
hello, Kitty
hello, Tom
exit cat named Tom
exit cat named Kitty

 

4、最后給出一個實例

使用上下文管理器實現redis分布式鎖

# -*- coding:utf-8 -*-
from __future__ import print_function
import redis
import time
import multiprocessing
from contextlib import contextmanager as _contextmanager
# 簡單創建redis的客戶端
r = redis.Redis(host='localhost', port=6379, db=0)

# 分布式鎖實現
# finally中驗證本線程是否獲得鎖, 是為了防止誤刪別的線程獲取的鎖
@_contextmanager
def dist_lock(client, key):
    dist_lock_key = 'lock:%s' % key
    is_acquire_lock = False
    try:
        is_acquire_lock = _acquire_lock(client, dist_lock_key)
        yield
    finally:
        if is_acquire_lock:
            _release_lock(client, dist_lock_key)


# 嘗試獲取鎖
# 成功: 返回True, 失敗: 拋出異常
# 使用set nx ex原語, 使得setnx和expire操作成為原子操作
def _acquire_lock(client, key):
    is_lock = r.set(key, 1, nx=True, ex=10)
    if not is_lock:
        raise Exception("already locked!")
    return is_lock


# 釋放鎖
# 簡單刪除key
# 如果刪除失敗, 鎖也會通過expire時間超時
def _release_lock(client, key):
    client.delete(key)


# 測試函數
# 獲取鎖成功, 打印成功, 並持有鎖3s
# 獲取鎖失敗, 直接打印
def func():
    while 1:
        try:
            with dist_lock(r, 'key'):
                print("*", end='')
                time.sleep(3)
        except Exception, ex:
            print('!', end='')


# 多進程啟動
# 這種模式下, 線程鎖無效, 可以驗證分布式鎖
process_list = list()
for i in range(2):
    process_list.append(multiprocessing.Process(target=func))
for process in process_list:
    process.start()
for process in process_list:
    process.join()