python-黑魔法-上下文管理器

python

python 上下文管理器

  1. with的作用是什么
  2. @contextlib.contextmanager与with有什么关系
  3. 上下文管理器使用案例
  4. 多个管理器嵌套 nested函数

问题一

with是一款语法糖 大致执行分为三个阶段

  1. enter(self): with进入前
  2. with 内部代码
  3. exit(self, type, value, traceback) with 退出后

问题二

  • 在理解 @contextlib.contextmanager内部先看下
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
#未使用上下文管理器，handle_query代码啰嗦
class Database(object):

    def __init__(self):
        self.connected = False

    def connect(self):
        self.connected = True

    def close(self):
        self.connected = False

    def query(self):
        if self.connected:
            return 'query data'
        else:
            raise ValueError('DB not connected ')

def handle_query():
    db = Database()
    db.connect()
    print 'handle --- ', db.query()
    db.close()

def main():
    handle_query()

if __name__ == '__main__':
    main()
  • 使用上下文管理器后
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
#handle_query 简单多了
import contextlib


class Database(object):

    def __init__(self):
        self.connected = False

    def connect(self):
        self.connected = True

    def close(self):
        self.connected = False

    def query(self):
        if self.connected:
            return 'query data'
        else:
            raise ValueError('DB not connected ')

@contextlib.contextmanager
def database():
    db = Database()
    try:
        if not db.connected:
            db.connect()
        yield db
    except Exception as e:
        db.close()


def handle_query():
    with database() as db:
        print('handle ---', db.query())

handle_query()
  • 分析下代码执行逻辑
  1. 进入@contextlib.contextmanager 实质是contextlib.contextmanager(database)
    1
    2
    3
    4
    5
    6
    #进入装饰器后又来一个装饰器wrap
    #简单说明下wrap装饰器的作用是将 要装饰的方法相关属性的指针指向装饰器，也就是将func的属性地址给 helper
    @wraps(func)
    def helper(*args, **kwds):
        return _GeneratorContextManager(func, args, kwds)
    return helper
  2. 执行到whit database() 的database()时候进入到return _GeneratorContextManager(func, args, kwds)，初始化 _GeneratorContextManager过程中 func（==database）协程初始化完成
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
class _GeneratorContextManager(ContextDecorator, AbstractContextManager):
    """Helper for @contextmanager decorator."""

    def __init__(self, func, args, kwds):
        self.gen = func(*args, **kwds)
        self.func, self.args, self.kwds = func, args, kwds
        # Issue 19330: ensure context manager instances have good docstrings
        doc = getattr(func, "__doc__", None)
        if doc is None:
            doc = type(self).__doc__
        self.__doc__ = doc
        # Unfortunately, this still doesn't provide good help output when
        # inspecting the created context manager instances, since pydoc
        # currently bypasses the instance docstring and shows the docstring
        # for the class instead.
        # See http://bugs.python.org/issue19404 for more details.

    def _recreate_cm(self):
        # _GCM instances are one-shot context managers, so the
        # CM must be recreated each time a decorated function is
        # called
        return self.__class__(self.func, self.args, self.kwds)

    def __enter__(self):
        try:
            return next(self.gen)
        except StopIteration:
            raise RuntimeError("generator didn't yield") from None

    def __exit__(self, type, value, traceback):
        if type is None:
            try:
                next(self.gen)
            except StopIteration:
                return False
            else:
                raise RuntimeError("generator didn't stop")
        else:
            if value is None:
                # Need to force instantiation so we can reliably
                # tell if we get the same exception back
                value = type()
            try:
                self.gen.throw(type, value, traceback)
            except StopIteration as exc:
                # Suppress StopIteration *unless* it's the same exception that
                # was passed to throw().  This prevents a StopIteration
                # raised inside the "with" statement from being suppressed.
                return exc is not value
            except RuntimeError as exc:
                # Don't re-raise the passed in exception. (issue27122)
                if exc is value:
                    return False
                # Likewise, avoid suppressing if a StopIteration exception
                # was passed to throw() and later wrapped into a RuntimeError
                # (see PEP 479).
                if type is StopIteration and exc.__cause__ is value:
                    return False
                raise
            except:
                # only re-raise if it's *not* the exception that was
                # passed to throw(), because __exit__() must not raise
                # an exception unless __exit__() itself failed.  But throw()
                # has to raise the exception to signal propagation, so this
                # fixes the impedance mismatch between the throw() protocol
                # and the __exit__() protocol.
                #
                if sys.exc_info()[1] is value:
                    return False
                raise
            raise RuntimeError("generator didn't stop after throw()")
  1. 触发__enter__ ，通过next(self.gen)启动协程
  2. 进入函数database yield db 挂起自己，释放控制权返回db
    1
    2
    3
    4
    5
    6
    7
    8
    def database():
        db = Database()
        try:
            if not db.connected:
                db.connect()
            yield db
        except Exception as e:
            db.close()
  3. 进入with内部 print(‘handle —‘, db.query())
  4. 触发__exit__ 继续回到协程 yeild db，关闭连接 退出with
1
2
问题二：
可以看到 with 是要与具有__enter__，__exit__的对象一起使用，而contextlib.contextmanager就是那个对象

问题三

1
2
3
4
5
6
7
8
9
10
11
12
13
#锁资源
@contextmanager
def locked(lock):
    lock.acquire()
    try:
        yield
    finally:
        lock.release()

with locked(myLock):
    #代码执行到这里时，myLock已经自动上锁
    pass
    #执行完后会，会自动释放锁
1
2
3
4
5
6
7
8
9
10
11
12
#文件资源
@contextmanager
def myopen(filename, mode="r"):
    f = open(filename,mode)
    try:
        yield f
    finally:
        f.close()

with myopen("test.txt") as f:
    for line in f:
        print(line)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
#数据库资源
@contextmanager
def transaction(db):
    db.begin()
    try：
        yield 
    except:
        db.rollback()
        raise
    else:
        db.commit()

with transaction(mydb):
    mydb.cursor.execute(sql)
    mydb.cursor.execute(sql)
    mydb.cursor.execute(sql)
    mydb.cursor.execute(sql)

问题4

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
from contextlib import contextmanager
from contextlib import nested
from contextlib import closing

@contextmanager
def my_context(name):
    print("enter")
    try:
        yield name
    finally:
        print("exit")

#使用nested函数来调用多个管理器
print("---------使用nested函数调用多个管理器-----------")
with nested(my_context("管理器一"), my_context("管理器二"),my_context("管理器三")) as (m1,m2,m3):
    print(m1)
    print(m2)
    print(m3)

#直接使用with来调用调用多个管理器
print("---------使用with调用多个管理器-----------")
with my_context("管理器一") as m1, my_context("管理器二") as m2, my_context("管理器三") as m3:
    print(m1)
    print(m2)
    print(m3)