| """ This module contains classes - NamedFileInTemporaryDirectory, TemporaryWorkingDirectory. | |
| These classes add extra features such as creating a named file in temporary directory and | |
| creating a context manager for the working directory which is also temporary. | |
| """ | |
| import os as _os | |
| from pathlib import Path | |
| from tempfile import TemporaryDirectory | |
| class NamedFileInTemporaryDirectory(object): | |
| def __init__(self, filename, mode="w+b", bufsize=-1, add_to_syspath=False, **kwds): | |
| """ | |
| Open a file named `filename` in a temporary directory. | |
| This context manager is preferred over `NamedTemporaryFile` in | |
| stdlib `tempfile` when one needs to reopen the file. | |
| Arguments `mode` and `bufsize` are passed to `open`. | |
| Rest of the arguments are passed to `TemporaryDirectory`. | |
| """ | |
| self._tmpdir = TemporaryDirectory(**kwds) | |
| path = Path(self._tmpdir.name) / filename | |
| encoding = None if "b" in mode else "utf-8" | |
| self.file = open(path, mode, bufsize, encoding=encoding) | |
| def cleanup(self): | |
| self.file.close() | |
| self._tmpdir.cleanup() | |
| __del__ = cleanup | |
| def __enter__(self): | |
| return self.file | |
| def __exit__(self, type, value, traceback): | |
| self.cleanup() | |
| class TemporaryWorkingDirectory(TemporaryDirectory): | |
| """ | |
| Creates a temporary directory and sets the cwd to that directory. | |
| Automatically reverts to previous cwd upon cleanup. | |
| Usage example: | |
| with TemporaryWorkingDirectory() as tmpdir: | |
| ... | |
| """ | |
| def __enter__(self): | |
| self.old_wd = Path.cwd() | |
| _os.chdir(self.name) | |
| return super(TemporaryWorkingDirectory, self).__enter__() | |
| def __exit__(self, exc, value, tb): | |
| _os.chdir(self.old_wd) | |
| return super(TemporaryWorkingDirectory, self).__exit__(exc, value, tb) | |