Python - with语句 - 写一个支持with方法调用的类

前言

使用语句with open('a.txt', 'r') as f:可以自动关闭文件句柄，使用语句with torcu.no_grad()可以不计算梯度。

我学Py的时候咋好像没学过open和no_grad是怎么实现的？

于是今日了解了一下。

实现一个my_with，计算语句执行用时

其实不难，写一个类，写两个函数__enter__（with代码块被调用之前执行）、__exit__（结束时候执行）。然后就可以了。

import time


class time_counter:    
    def __enter__(self):
        self.start = time.time()
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.end = time.time()
        print(f"time consume: {self.end - self.start}s")


with time_counter():
    for i in range(int(1e7)):
        i += 1

运行结果：

1	`time consume: 0.7823185920715332s`

那“as”语句呢？

只需要给__enter__函数加个返回值：

import time


class time_counter:    
    def __enter__(self) -> int:
        self.start = time.time()
        return 1
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.end = time.time()
        print(f"time consume: {self.end - self.start}s")


with time_counter() as a:
    print(a)
    for i in range(int(1e7)):
        i += 1