如何为类实现类似with语句的功能？

Question

抱歉我起的标题不太好，我想不出更好的了。如果你有好的标题建议，欢迎告诉我。

我想做一个可以处理HDF5文件的接口，这个接口支持多进程同时访问文件，并且通过文件锁来管理。这个模块的使用环境是在一个Linux集群中，大家通过NFS共享一个磁盘。我的目标是让多个不同主机上的并行进程能够同时访问同一个文件。

我希望能通过一个包装类来实现锁定功能，这个包装类是基于h5py.File类的。（h5py已经支持线程级别的并发，但底层的HDF5库并不是线程安全的。）

如果我能做一些类似这样的事情就太好了：

class LockedH5File(object):
    def __init__(self, path, ...):
        ...
        with h5py.File(path, 'r+') as h5handle:
            fcntl.flock(fcntl.LOCK_EX)
            yield h5handle
        # (method returns)

我知道上面的代码是错的，但我希望它能传达主要的想法：也就是说，使用LockedH5File('/path/to/file')可以给客户端代码返回一个打开的文件句柄，客户端可以在这个句柄上进行各种读写操作。当这个句柄超出作用域时，它的析构函数会关闭这个句柄，从而释放锁。

促使我这样设计的目标有两个：

1. 将句柄的创建（由库代码完成）和后续对句柄的操作（由客户端代码请求）分开，

2. 确保无论在执行过程中发生什么情况（比如异常、未处理的信号、Python内部错误），句柄都能被关闭，锁也能被释放。

我该如何在Python中实现这个效果呢？

谢谢！

Answer 1

好的，首先我们来聊聊上下文管理器和 with 语句。一般来说，在Python中，析构函数（也就是对象被销毁时调用的函数）并不一定会执行，所以你不能指望它们能完成所有的清理工作。你应该把它们当作一种保险措施来使用。要使用上下文管理器，你需要提供 __enter__ 和 __exit__ 这两个方法，然后可以像这样使用它：

with LockedFile(...) as fp:
    # ...

Answer 2

可以在 with 语句中使用的对象被称为上下文管理器，它们有一个简单的接口。上下文管理器需要提供两个方法：一个是 __enter__ 方法，这个方法不需要任何参数，可以返回任何东西（这个返回值会被赋值给 as 部分中的变量）；另一个是 __exit__ 方法，这个方法需要三个参数（这些参数会用 sys.exc_info() 的结果填充），并且返回一个非零值来表示异常已经被处理。你的例子可能看起来像这样：

class LockedH5File(object):
    def __init__(self, path, ...):
        self.path = path

    def __enter__(self):
        self.h5handle = h5handle = h5py.File(self.path, 'r+')
        fcntl.flock(fcntl.LOCK_EX)
        return h5handle

    def __exit__(self, exc_type, exc_info, exc_tb):
        self.h5handle.close()

Answer 3

要让这个工作正常，你的类需要遵循上下文管理器协议。另外，你也可以使用contextlib.contextmanager这个装饰器来写一个生成器函数。

你的类大概可以长成这样（关于h5py的用法细节可能是错的）：

class LockedH5File(object):
    def __init__(self, path, ...):
        self.h5file = h5py.File(path, 'r+')
    def __enter__(self):
        fcntl.flock(fcntl.LOCK_EX)
        return self.h5file
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.h5file.close()