You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
2 years ago
|
""" 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)
|