diff options
Diffstat (limited to 'Lib/tempfile.py')
-rw-r--r-- | Lib/tempfile.py | 65 |
1 files changed, 64 insertions, 1 deletions
diff --git a/Lib/tempfile.py b/Lib/tempfile.py index 049cdaa2c2..699fd0a9f9 100644 --- a/Lib/tempfile.py +++ b/Lib/tempfile.py @@ -19,7 +19,7 @@ This module also provides some data items to the user: __all__ = [ "NamedTemporaryFile", "TemporaryFile", # high level safe interfaces - "SpooledTemporaryFile", + "SpooledTemporaryFile", "TemporaryDirectory", "mkstemp", "mkdtemp", # low level safe interfaces "mktemp", # deprecated unsafe interface "TMP_MAX", "gettempprefix", # constants @@ -613,3 +613,66 @@ class SpooledTemporaryFile: def xreadlines(self, *args): return self._file.xreadlines(*args) + + +class TemporaryDirectory(object): + """Create and return a temporary directory. This has the same + behavior as mkdtemp but can be used as a context manager. For + example: + + with TemporaryDirectory() as tmpdir: + ... + + Upon exiting the context, the directory and everthing contained + in it are removed. + """ + + def __init__(self, suffix="", prefix=template, dir=None): + self.name = mkdtemp(suffix, prefix, dir) + self._closed = False + + def __enter__(self): + return self.name + + def cleanup(self): + if not self._closed: + self._rmtree(self.name) + self._closed = True + + def __exit__(self, exc, value, tb): + self.cleanup() + + __del__ = cleanup + + + # XXX (ncoghlan): The following code attempts to make + # this class tolerant of the module nulling out process + # that happens during CPython interpreter shutdown + # Alas, it doesn't actually manage it. See issue #10188 + _listdir = staticmethod(_os.listdir) + _path_join = staticmethod(_os.path.join) + _isdir = staticmethod(_os.path.isdir) + _remove = staticmethod(_os.remove) + _rmdir = staticmethod(_os.rmdir) + _os_error = _os.error + + def _rmtree(self, path): + # Essentially a stripped down version of shutil.rmtree. We can't + # use globals because they may be None'ed out at shutdown. + for name in self._listdir(path): + fullname = self._path_join(path, name) + try: + isdir = self._isdir(fullname) + except self._os_error: + isdir = False + if isdir: + self._rmtree(fullname) + else: + try: + self._remove(fullname) + except self._os_error: + pass + try: + self._rmdir(path) + except self._os_error: + pass |