| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162 |
- """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 io import BufferedWriter
- from pathlib import Path
- from tempfile import TemporaryDirectory
- from types import TracebackType
- from typing import Optional, Type
- class NamedFileInTemporaryDirectory:
- def __init__(self, filename: str, mode: str, bufsize: int=-1, add_to_syspath: bool=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) -> BufferedWriter:
- return self.file
- def __exit__(self, type: Optional[Type[BaseException]], value: Optional[BaseException], traceback: Optional[TracebackType]):
- 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) -> str:
- self.old_wd = Path.cwd()
- _os.chdir(self.name)
- return super(TemporaryWorkingDirectory, self).__enter__()
- def __exit__(self, exc: Optional[Type[BaseException]], value: Optional[BaseException], tb: Optional[TracebackType]) -> None:
- _os.chdir(self.old_wd)
- return super(TemporaryWorkingDirectory, self).__exit__(exc, value, tb)
|
