python 魔术方法上下文管理（__enter__,__exit__）

上下文管理

文件IO操作可以对文件对象使用上下文管理，使用with……as语法。

with open("test") as f:
    pass

仿照上例写一个自己的类，实现上下文管理。

class Point:
    pass

with Point() as p:
    pass

结果为：
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-1-77f0ac1b02c7> in <module>
      2     pass
      3 
----> 4 with Point() as p:
      5     pass

AttributeError: __enter__

提示属性错误，没有__exit__，看了需要这个属性。

上下文管理对象

当一个对象同时实现了__enter__()和__exit__()方法，它就属于上下文管理对象。上下文管理器的主要原理是你的代码会放到 with 语句块中执行。 当出现 with 语句的时候，对象的 __enter__() 方法被触发， 它返回的值(如果有的话)会被赋值给 as 声明的变量。然后，with 语句块里面的代码开始执行。 最后，__exit__() 方法被触发进行清理工作。

__enter__()：进入于此相关的上下文。如果存在该方法，with语法会把该方法的返回值作为绑定到as子句中指定的变量上。

__exit__：退出与此对象相关的上下文。

class Point:
    def __init__(self):
        print("init")
        
    def __enter__(self):#进去的时候做的事
        print("enter")
        
    def __exit__(self,exc_type,exc_val,exc_tb):#离开的时候做的事情
        print("exit")
        
with Point() as f:
    print("do sth.")

结果为：
init
enter
do sth.
exit

实例化对象的时候，并不会调用enter，而是进入with语句块调用__enter__方法，然后执行语句块，最后离开with语句块的时候，调用__exit__方法。

with可以开启一个上下文运行环境，在执行前做一些准备工作，执行后做一些收尾工作。

上下文管理的安全性

看看异常对上下文管理的影响。

class Point:
    def __init__(self):
        print("init")
        
    def __enter__(self):
        print("enter")
        
    def __exit__(self,exc_type,exc_val,exc_tb):
        print("exit")
        
with Point() as f:
    raise Exception("error")
    print("do sth.")

结果为：
init
enter
exit
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-4-28a514a324bc> in <module>
     10 
     11 with Point() as f:
---> 12     raise Exception("er")
     13     print("do sth.")

Exception: er

可以看到在enter和exit照样执行，上下文管理是安全的。

极端的例子

调用sys.exit()，它会退出当前解释器。

打开Python解释器，在里面敲入sys.exit()，窗口直接关闭了，也就是说碰到这一句，Python运行环境直接退出了。

import sys

class Point:
    def __init__(self):
        print("init")
        
    def __enter__(self):
        print("enter")
        
    def __exit__(self,exc_type,exc_val,exc_tb):
        print("exit")
        
with Point() as f:
    sys.exit(-100)
    print("do sth.")

print("outer")

结果为：
init
enter
exit
An exception has occurred, use %tb to see the full traceback.

SystemExit: -100


d:Anaconda3libsite-packagesIPythoncoreinteractiveshell.py:3304: UserWarning: To exit: use 'exit', 'quit', or Ctrl-D.
  warn("To exit: use 'exit', 'quit', or Ctrl-D.", stacklevel=1)

从执行结果来看，依然执行了__exit__函数，哪怕是退出Python运行环境，说明上下文管理很安全。

class Point:
    def __init__(self):
        print("init")
        
    def __enter__(self):
        print("enter"+self.__class__.__name__)
        return self
    
    def __exit__(self,exc_type,exc_val,exc_tb):
        print("exit"+self.__class__.__name__)
        print(exc_type)
        print(exc_val)
        print(exc_tb)
        return 1
        
p = Point()
with p as f:
    raise Exception("error123")
    
    print(f==p)
    print(f is p)
    print(f)
    print(p)
    

结果为：
init
enterPoint
exitPoint
<class 'Exception'>
error123
<traceback object at 0x03AE2FD0>

class Point:
    def __init__(self):
        print("init")
        
    def __enter__(self):
        print("enter"+self.__class__.__name__)
        return self
    
    def __exit__(self,exc_type,exc_val,exc_tb):
        print("exit"+self.__class__.__name__)
        print(exc_type)
        print(exc_val)
        print(exc_tb)
        return 0
        
p = Point()
with p as f:
    raise Exception("error123")
    
    print(f==p)
    print(f is p)
    print(f)
    print(p)
    

结果为：
init
enterPoint
exitPoint
<class 'Exception'>
error123
<traceback object at 0x03D30F58>
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-6-628d15869af5> in <module>
     16 p = Point()
     17 with p as f:
---> 18     raise Exception("error123")
     19 
     20     print(f==p)

Exception: error123

由上面的例子可以看到，返回值很重要，一个返回值可以压制异常，是外面管还是里面管，返回等效true的话就是里面管，返回false的话就是外面管。所以这两个函数的返回值都很重要。

with语句

class Point:
    def __init__(self):
        print("init")
        
    def __enter__(self):
        print("enter")
        
    def __exit__(self,exc_type,exc_val,exc_tb):print("exit")
        

p = Point()
with p as f:
    print(p == f)#为什么不相等
    print("do sth.")

结果为：

init
enter
False
do sth.
exit

问题在于__enter__方法，它将自己的返回值赋给f,修改上例。

class Point:
    def __init__(self):
        print("init")
        
    def __enter__(self):
        print("enter")
        return self
        
    def __exit__(self,exc_type,exc_val,exc_tb):
        print("exit")

p = Point()
with p as f:
    print(p == f)#为什么不相等
    print("do sth.")

结果为：
init
enter
True
do sth.
exit

__enter__方法返回值就是上下文中使用的对象，with语法会把它的返回值赋给as子句的变量。

__enter__方法和__exit__方法的参数

__enter__方法没有其他参数。

__exit__方法有三个参数。

__exit__(self,exc_type,exc_value,traceback)

这三个参数都是与异常相关，如果该上下文退出时没有异常，这3个参数都为None。如果有异常，参数意义如下：

exc_type:异常类型

exc_vaule:异常的值

traceback：异常的追踪信息。

__exit__方法返回一个等效true的值，则压制异常，否则，继续抛出异常。

class Point:
    def __init__(self):
        print("init")
        
    def __enter__(self):
        print("enter")
        
    def __exit__(self,exc_type,exc_val,exc_tb):
        print(exc_type)
        print(exc_val)
        print(exc_tb)
        print("exit")
        return "abc"

p = Point()
with p as f:
    raise Exception("new error")
    print("do sth.")

print("outer")

结果为：
init
enter
<class 'Exception'>
new error
<traceback object at 0x0000000005A7A088>
exit
outer

练习

为加法函数计时

方法1、使用装饰器显示该函数的执行时长

方法2、商用上下文管理方法来显示该函数的执行时长。

import time

def add(x,y):
    time.sleep(2)
    return x+y

装饰器实现

import time
import datetime
from functools import wraps

def timeit(fn):
    @wraps(fn)
    def wrapper(*args,**kwargs):
        start = datetime.datetime.now()
        ret = fn(*args,**kwargs)
        delta = (datetime.datetime.now()-start).total_seconds()
        print("{} took {}s".format(fn.__name__,delta))
        return ret
    return wrapper

@timeit
def add(x,y):
    time.sleep(2)
    return x+y

print(add(4,5))

结果为：
add took 2.0s
9

上下文管理

import time
import datetime
from functools import wraps

def add(x,y):
    time.sleep(2)
    return x+y

class Timeit():
    def __init__(self,fn):
        self.fn = fn
        
    def __enter__(self):
        self.start =datetime.datetime.now()
        
    def __exit__(self,exc_type,exc_val,exc_tb):
        delta = (datetime.datetime.now()-self.start).total_seconds()
        print("{} took {}s".format(self.fn.__name__,delta))
        
with Timeit(add) as fn:
    print(add(4,7))

结果为：
11
add took 2.001s

另一种实现，使用可调用对象实现。

import time
import datetime
from functools import wraps

def add(x,y):
    time.sleep(2)
    return x+y

class Timeit():
    def __init__(self,fn):
        self.fn = fn
        
    def __enter__(self):
        self.start =datetime.datetime.now()
        return self
        
    def __exit__(self,exc_type,exc_val,exc_tb):
        delta = (datetime.datetime.now()-self.start).total_seconds()
        print("{} took {}s".format(self.fn.__name__,delta))
        
    def __call__(self,x,y):
        print(x,y)
        return self.fn(x,y)
    
with Timeit(add) as timeitobj:
    print(timeitobj(4,5))

根据上面的代码，能不能把类当做装饰器用？

import time
import datetime
from functools import wraps

class Timeit:
    """ this is a class"""
    def __init__(self,fn):
        self.fn = fn
        #把函数对象的文档字符串赋值给类
        self.__doc__ = fn.__doc__
        #update_wrapper(self,fn)
        
    def __enter__(self):
        self.start = datetime.datetime.now()
        return self
    
    def __exit__(self,exc_type,exc_val,exc_tb):
        delta = (datetime.datetime.now()-self.start).total_seconds()
        print("{} took {}s".format(self.fn.__name__,delta))
        
    def __call__(self,*args,**kwargs):
        self.start = datetime.datetime.now()
        ret = self.fn(*args,**kwargs)
        delta = (datetime.datetime.now()-self.start).total_seconds()
        print("{} took {}s".format(self.fn.__name__,delta))
        return ret
    
@Timeit
def add(x,y):
    """ this is add function"""
    time.sleep(2)
    return x+y

print(add(10,5))
print(add.__doc__)

print(Timeit(add).__doc__)

结果为：
add took 2.018236s
15
 this is add function
 this is add function

思考，如何解决文档字符串问题？

方法一：直接修改__doc__

class Timeit:
    def __init__(self,fn = None):
        self.fn = fn
        #把函数对象的文档字符串赋给类
        self.__doc__ = fn.__doc__

方法二：使用functools.wraps函数

import time
import datetime
from functools import wraps

class Timeit:
    """ this is a class"""
    def __init__(self,fn):
        self.fn = fn
        #把函数对象的文档字符串赋值给类
        self.__doc__ = fn.__doc__
        #update_wrapper(self,fn)
        
    def __enter__(self):
        self.start = datetime.datetime.now()
        return self
    
    def __exit__(self,exc_type,exc_val,exc_tb):
        delta = (datetime.datetime.now()-self.start).total_seconds()
        print("{} took {}s".format(self.fn.__name__,delta))
        
    def __call__(self,*args,**kwargs):
        self.start = datetime.datetime.now()
        ret = self.fn(*args,**kwargs)
        delta = (datetime.datetime.now()-self.start).total_seconds()
        print("{} took {}s".format(self.fn.__name__,delta))
        return ret
    
@Timeit
def add(x,y):
    """ this is add function"""
    time.sleep(2)
    return x+y

print(add(10,5))
print(add.__doc__)

print(Timeit(add).__doc__)
        

上面的类即可以用在上下文管理，又可以用作装饰器，初始化init函数不写return不代表返回是none，它返回的是实例。

上下文应用场景

1.增强功能

在代码执行的前后增加代码，以增强其功能，类似装饰器的功能。

2.资源管理

打开了资源需要关闭，例如文件对象，网络连接，数据库连接

3.权限验证

在执行代码之前，做权限的验证，在__enter__中处理。

contexlib.contextanmger

contexlib.contextanmger是一个装饰器实现上下文管理， 装饰一个函数，而不用像类一样实现__enter__和__exit__方法。

对下面的函数有要求，必须有yield，也就是这个函数必须返回一个生成器，且只有yield一个值，也就是这个装饰器接收一个生成器对象作为参数。

import contextlib

@contextlib.contextmanager
def foo():
    print("enter")#相当于__enter__()
    yield #yield 5,yield的值只能有一个，作为__enter__方法的返回值
    print("exit") #相当于__exit__()
    
with foo() as f:
    #raise Exception()
    print(f)

结果为：

enter
None
exit

import contextlib

@contexlib.contextanmger
def foo():
    print("enter")#相当于__enter__()
    yield #yield 5,yield的值只能有一个，作为__enter__方法的返回值
    print("exit") #相当于__exit__()
    
with foo() as f:
    #raise Exception()
    print(f)

结果为：

enter
None
exit

f接收yield语句的返回值

上面的程序看似不错，但是，增加一个异常试一试，发现不能保证exit的执行，怎么办？

import contextlib

@contextlib.contextmanager
def foo():
    print("enter")#相当于__enter__()
    try:
        yield #yield 5,yield的值只能有一个，作为__enter__方法的返回值
    finally:
        print("exit") 
    
with foo() as f:
    raise Exception()
    print(f)

结果为：
enter
exit
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-37-36a79d2a42b7> in <module>
     10 
     11 with foo() as f:
---> 12     raise Exception()
     13     print(f)

Exception: 

上面这么做有什么意义呢？

当yield发生处为生成器函数增加了上下文管理，这是为函数增加上下文机制的方式。

• 把yield之前的当做__enter__方法执行
• 把yield之后的当做__exit__方法执行
• 把yield的值作为__enter__的返回值

import contextlib
import time
import datetime

@contextlib.contextmanager
def add(x,y):#为生成器函数增加了上下文管理
    start= datetime.datetime.now()
    try:
        yield x+y#yield5,yield的值只能有一个，作为__enter__方法的返回值
    finally:
        delta = (datetime.datetime.now()-start).total_seconds()

with add(4,5) as f:
    #raise Exception()
    time.sleep(2)
    print(f)

结果为：
9

总结：

如果业务逻辑简单可以使用函数加contextlib.contextmanager装饰器方式，如果业务复杂，用类的方式加__enter__和__exit__方法方便。