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