Python--魔术方法--上下文管理

上下文管理

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

with open('test') as f:
	pass

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

class Point：
	pass

with Point() as p: # AttributeError: __exit__
	pass

提示属性错误，没有__exit__，看了需要这个属性
某些版本会显示没有__enter__

上下文管理对象

当一个对象同时实现了__enter()__和__exit__()方法，它就属于上席文管理的对象

方法	意义
__enter__	进入与此对象相关的上下文。如果存在该方法，with语法会把该方法的返回值作为绑定到as子句中指定的变量上
__exit__	退出与此对象相关的上下文

import time

class Point():
    def __init__(self):
        print('init =========')
        time.sleep(1)
        print('init over =====')

    def __enter__(self):
        print('enter ========')

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('exit =========')

with Point() as p:
    print('in with =========')
    time.sleep(2)
    print('with over ========')

print('======end======')

# 执行结果
init =========
init over =====
enter ========
in with =========
with over ========
exit =========
======end======

• 实例化对象的时候，并不会调用enter，进入with语句块调用__enter__方法，然后执行语句体，最后离开with语句块的时候，调用__exit__方法
• with可以开启一个上下文运行环境，在执行前做一些准备工作，执行后做一些首位工作。
• 注意，with并不开启一个新的作用域

上下文管理的安全性

• 异常对上下文的影响

import time

class Point():
    def __init__(self):
        print('init =========')
        time.sleep(1)
        print('init over =====')

    def __enter__(self):
        print('enter ========')

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('exit =========')

with Point() as p:
    print('in with =========')
    raise Exception('error')
    time.sleep(2)
    print('with over ========')

print('======end======')

执行结果

虽然有异常弹出，但是enter和exit照样执行，上下文管理是安全的

• 极端的例子
  调用sys.exit()，它会退出当前解释器
  打开python解释器，在里面输入 sys.exit()，窗口直接关闭，也就是说碰到这一句，Python运行环境直接退出

import time

class Point():
    def __init__(self):
        print('init =========')
        time.sleep(1)
        print('init over =====')

    def __enter__(self):
        print('enter ========')

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('exit =========')

with Point() as p:
    print('in with =========')
    import sys
    sys.exit(1)
    time.sleep(2)
    print('with over ========')

print('======end======')

执行结果

init =========
init over =====
enter ========
in with =========
exit =========

Process finished with exit code 1

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

with语句

class Point():
    def __init__(self):
        print('init')

    def __enter__(self):
        print('enter')

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('exit')

f = open('t3.py')
with f as p:
    print(f)
    print(p)
    print(1,f is p) # 打印什么
    print(2,f == p) # 打印什么

p = Point()
with p as f:
    print('in with-----')
    print(3,p == f)
    print('with over')

print('======end======')

执行结果

<_io.TextIOWrapper name='t3.py' mode='r' encoding='cp936'>
<_io.TextIOWrapper name='t3.py' mode='r' encoding='cp936'>
1 True
2 True
init
enter
in with-----
3 False
with over
exit
======end======

问题在于__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('in with-----')
    print(p == f)
    print('with over')

print('======end======')

执行结果

init
enter
in with-----
True
with over
exit
======end======

with语句，会调用with后的对象的__enter__方法，如果有as，则将该方法的返回值赋给as子句的变量
上例，可以等价为f = p.__enter__()

方法的参数

__enter__方法 没有其他参数
__exit__方法有3个参数
__exit__(self, exc_type, exc_value, traceback)

• 这三个参数都与异常有关。
• 如果该上下文退出时没有异常，这3个参数都为None
• 如果有异常，参数意义如下：
  exc_type，异常类型
  exc_value，异常的值
  traceback，异常的追踪信息
• __exit__方法返回一个等效True的值，则压制异常；否则，继续抛出异常

class Point():
    def __init__(self):
        print('init')

    def __enter__(self):
        print('enter')
        return self # 增加返回值

    def __exit__(self, exc_type, exc_val, exc_tb):
        print(1, exc_type)
        print(2, exc_val)
        print(3, exc_tb)
        print('exit')
        return 'abc' # return 值恒等为True，压制报错
        # return None # 0 # []

p = Point()
with p as f:
    print('in with-----')
    raise Exception('Error')
    print('with over')

print('======end======')

执行结果

init
enter
in with-----
1 <class 'Exception'>
2 Error
3 <traceback object at 0x00000259B3BFDB88>
exit
======end======

练习

为加法函数计时
方法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))

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

# 执行结果
add took 2.000913s
9

上下文实现

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

class TimeIt:
    def __init__(self, fn):
        self.fn = fn

    def __enter__(self):
        self.start = datetime.datetime.now()
        return self.fn

    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))
   
# 执行结果
add took 2.000563s
11
add took 2.000563s

可调用对象实现

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

class TimeIt:
    def __init__(self, fn):
        self.fn = fn

    def __enter__(self):
        self.start = datetime.datetime.now()
        return self.fn

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

# 执行结果
add took 2.000302s
11
add took 2.000302s

将上面代码的类当装饰器使用

import time
import datetime
from functools import wraps

class TimeIt:
    def __init__(self, fn):
        self.fn = fn

    def __enter__(self):
        self.start = datetime.datetime.now()
        return self.fn

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

    def __call__(self, *args, **kwargs):
        self.start = datetime.datetime.now()
        ret = self.fn(*args, **kwargs)
        self.delta = (datetime.datetime.now() - self.start).total_seconds()
        print('{} took {}s'.format(self.fn.__name__, self.delta))
        return ret

@TimeIt
def add(x, y):
    """This is add function."""
    time.sleep(2)
    return x + y

add(4, 5)
print(add.__doc__)
# print(add.__name__) # 异常，没有此属性

# 执行结果
add took 2.000581s
None

解决文档字符串的问题

• 方法一
  直接修改__doc__

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

• 方法2
  使用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)
        wraps(fn)(self)

    def __enter__(self):
        self.start = datetime.datetime.now()
        return self.fn

    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)
        self.delta = (datetime.datetime.now() - self.start).total_seconds()
        print('{} took {}s'.format(self.fn.__name__, self.delta))
        return ret

@TimeIt
def add(x, y):
    """This is add function."""
    time.sleep(2)
    return x + y

print(add(4, 5))
print(add.__doc__)
print(TimeIt(add).__doc__)

# 执行结果
add took 2.000641s
9
This is add function.
This is add function.

上面的类即可以用在上下文管理，又可以用作装饰器

上下文应用场景

1. 增强功能
   在代码执行的前后增加代码，以增强其功能。类似装饰器的功能
2. 资源管理
   打开了资源需要关闭，例如文件对象、网络连接、数据库连接等
3. 权限验证
   在执行代码之前，做权限的验证，在__enter__中处理

contextlib.contextmanager

• 它是一个装饰器实现上下文管理，装饰一个函数，而不用像类一样实现__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

as后的f接受yield语句的返回值

• 上面程序如果打开 raise Exception()语句，则print('exit')不会执行
• 解决办法：增加try finally

import contextlib

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

with foo() as f:
    raise Exception()
    print(f)

执行结果

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

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

import contextlib
import datetime
import time

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

with foo(4, 5) as f:
    # raise Exception()
    print(f)

# 执行结果
9
2.000907

总结

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