上下文管理器

何谓上下文管理器

说到上下文，相信大家都还记得，在中学语文考试的阅读理解题上，我们经常会看到“联系上下文，解释xxx的意思”，这里的“上下文”指的是语言环境。

那么对应到代码中，其实也是一个意思，即程序所执行的环境状态。

那么管理器又怎么说呢？在我们写代码的时候，我们会把操作放在一个代码块中（比如读写文件），这样执行代码块时就可以保持一种运行状态，而当离开这个代码块时，就结束当前状态，去执行另一个操作，如果在这个代码块的运行状态还没结束时，此时又“越界”去做了另一个操作，上下文管理器就会做出相应的处理。

在我们做IO操作时（比如文件读写、数据库连接断开等操作），每一个操作都是会占用系统资源的，而且系统资源有限，如果使用这些资源后却一直不释放，那么久容易造成资源泄露，导致系统运行缓慢，甚至崩溃。

比如这样一个例子：

1
2
3

for x in range(10000000): 
    f = open('test.txt', 'w')
    f.write('python')

打开test.txt这个文件并写入一句话，执行10000000次，但是每次其实是没有去释放资源的（即关闭文件），这样就很容易造成资源泄露，导致报错OSError。

而更合理的做法应该是，每次打开写入完成后，就关闭文件去释放资源，所以为了解决这个问题，通过上下文管理器可以帮助我们使用完资源后去自动释放。

1
2
3

for x in range(10000000):
    with open('test.txt','w') as f :
        f.write('python')

那么我们想要自己实现一个上下文管理器，该怎么做呢？

实现上下文管理器

实现上下文管理器的方法，有两种，一种是通过魔术方法实现，一种是基于生成器实现。

魔术方法实现

上下文管理器的实现，涉及到__enter__和__exit__两个魔术方法，我们自己来定义一个HandleFile类，来实现文件打开和自动关闭。

class HandleFile:
    def __init__(self, name, mode):
        print('调用__init__')
        self.name = name
        self.mode = mode
        self.file = None

    def __enter__(self):
        print('调用__enter__')
        self.file = open(self.name, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('调用__exit__')
        self.file.close()


with HandleFile('test.txt', 'w') as f:
    print('在文件中写入')
    f.write('python')
# 输出
调用__init__
调用__enter__
在文件中写入
调用__exit__

通过with 关键字激活上下文管理HandleFile后，HandleFile('test.txt', 'w')先初始化完成，然后自动执行__enter__ 方法进入上下环境，__enter__ 方法的返回值赋值给变量f，f在调用write方法完成写操作后，自动触发__exit__来关闭文件。

这两个魔术方法的详细使用，可以查看《魔术方法总结》一文中的对应小节，这里不再赘述。

基于生成器实现

如果不想通过一个类来实现一个简单的上下文管理器，python提供了一个装饰器contextlib.contextmanager，通过它可以将一个函数变为上下文管理器。

我们来实现一个和魔术方法等价的上下文管理器：

from contextlib import contextmanager


@contextmanager
def handle_file(name, mode):
    try: # 相当于__enter__
        f = open(name, mode)
        yield f
    except Exception:
        print('出错了')
    finally: # 相当于__exit__
        f.close()


with handle_file('test.txt', 'w') as f:
    f.write('python666')
    print(1/0)
# 输出
出错了

在被装饰的函数中，必须有yiled关键字，yiled之前的代码就相当于__enter__，完成打开文件并将返回值传给f，yiled之后的关键字就相当于__exit__，等with代码块中的代码执行完成后，就去关闭文件。此外，通过try-except 实现了上下管理器的异常处理。所以，这个方法和魔术方法实现的上下文管理器，是完全等价的。

在工作中的应用

数据库连接断开

在做接口自动化时，数据库的连接断开，也可以通过上下文管理器来简化：

import pymysql


class HandleMysql:
    def __init__(self, hostname, port):
        self.hostname = hostname
        self.port = port
        self.connection = None

    def __enter__(self):
        self.connection = pymysql.connect(self.hostname, self.port, charset="utf8")
        return self.connection

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.connection.close()


with HandleMysql('localhost', '80') as db:
    print('执行数据库的增删查改操作')

allure.step()

在allure中提供了一个step()方法，通过这个方法，可以将测试步骤展现在测试报告中。step 有两种用法，一种通过装饰器调用，一种则是通过上下文管理器使用。

# @allure.step('step1:注册') 方法一
def case_register(self, data):
    """
    注册业务场景
    :param data:
    :return:
    """
    with allure.step('step1:注册'): # 方法二
        data = self.template(data, {'mobile_phone': self.random_phone()})
        res = self.register_api(**data).json()
        return res

我们点进step方法查看源码：

def step(title):
    if callable(title):
        return StepContext(title.__name__, {})(title)
    else:
        return StepContext(title, {})

可以看到返回的是StepContext()对象，继续点进去：

class StepContext:

    def __init__(self, title, params):
        self.title = title
        self.params = params
        self.uuid = uuid4()

    def __enter__(self):
        plugin_manager.hook.start_step(uuid=self.uuid, title=self.title, params=self.params)

    def __exit__(self, exc_type, exc_val, exc_tb):
        plugin_manager.hook.stop_step(uuid=self.uuid, title=self.title, exc_type=exc_type, exc_val=exc_val,
                                      exc_tb=exc_tb)

    def __call__(self, func):
        @wraps(func)
        def impl(*a, **kw):
            __tracebackhide__ = True
            params = func_parameters(func, *a, **kw)
            args = list(map(lambda x: represent(x), a))
            with StepContext(self.title.format(*args, **params), params):
                return func(*a, **kw)
        return impl