@contextlib.contextmanager is awesome to easily create contextmanagers without boilerplate. However, creating contextmanager for classes is still painful/boilerplate, expecially when dealing with nested contextmanagers.
It would be great to extend the yield-based syntax of
@contextlib.contextmanager to classes (e.g. through a
Here is an example how this would look on a (simplified) real use-case I had.
@contextmanager @dataclass class File: path: str mode: str = 'r' def __contextmanager__(self): with tf.io.gfile.GFile(self.path, self.mode) as f: with h5py.File(f, self.mode) as h5_f: yield h5_f ... # File has other methods with File('/path/to/file.txt') as f: data = f['dataset']
__exit__ would be much more verbose/ugly. One would have to save the
self._gfile_context attributes. It’s also not trivial at all to correctly implement
__exit__ in case
h5py.File suppress the exceptions (
The implementation could be a simple extension of
contextlib.contextmanager. Here is a proof of concept:
def contextmanager(cls): cm: Optional[ContextManager[_T]] = None def __enter__(self): nonlocal cm cm = self.__contextmanager__() return cm.__enter__() def __exit__(self, exc_type, exc_value, traceback): return cm.__exit__(exc_type, exc_value, traceback) cls.__enter__ = __enter__ cls.__contextmanager__ = contextlib.contextmanager(cls.__contextmanager__) cls.__exit__ = __exit__ return cls
Note: This implementation also works with inheritance
@contextmanager class FileWrapper(File): def __contextmanager__(self): with super().__contextmanager__() as f: yield f