|
|
- """TemporaryDirectory class, copied from Python 3
-
- NamedFileInTemporaryDirectory and TemporaryWorkingDirectory from IPython, which
- uses the 3-clause BSD license.
- """
- from __future__ import print_function
-
- import os as _os
- import warnings as _warnings
- import sys as _sys
-
- # This code should only be used in Python versions < 3.2, since after that we
- # can rely on the stdlib itself.
- try:
- from tempfile import TemporaryDirectory
-
- except ImportError:
- from tempfile import mkdtemp, template
-
- 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 everything 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, _warn=False):
- if self.name and not self._closed:
- try:
- self._rmtree(self.name)
- except (TypeError, AttributeError) as ex:
- # Issue #10188: Emit a warning on stderr
- # if the directory could not be cleaned
- # up due to missing globals
- if "None" not in str(ex):
- raise
- print("ERROR: {!r} while cleaning up {!r}".format(ex, self,),
- file=_sys.stderr)
- return
- self._closed = True
- if _warn:
- self._warn("Implicitly cleaning up {!r}".format(self),
- Warning)
-
- def __exit__(self, exc, value, tb):
- self.cleanup()
-
- def __del__(self):
- # Issue a ResourceWarning if implicit cleanup needed
- self.cleanup(_warn=True)
-
-
- # 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
- _warn = _warnings.warn
-
- 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
-
-
- class NamedFileInTemporaryDirectory(object):
- """Open a file named `filename` in a temporary directory.
-
- This context manager is preferred over :class:`tempfile.NamedTemporaryFile`
- when one needs to reopen the file, because on Windows only one handle on a
- file can be open at a time. You can close the returned handle explicitly
- inside the context without deleting the file, and the context manager will
- delete the whole directory when it exits.
-
- Arguments `mode` and `bufsize` are passed to `open`.
- Rest of the arguments are passed to `TemporaryDirectory`.
-
- Usage example::
-
- with NamedFileInTemporaryDirectory('myfile', 'wb') as f:
- f.write('stuff')
- f.close()
- # You can now pass f.name to things that will re-open the file
- """
- def __init__(self, filename, mode='w+b', bufsize=-1, **kwds):
- self._tmpdir = TemporaryDirectory(**kwds)
- path = _os.path.join(self._tmpdir.name, filename)
- self.file = open(path, mode, bufsize)
-
- 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 = _os.getcwd()
- _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)
-
|