You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

154 lines
5.1 KiB

4 years ago
  1. """TemporaryDirectory class, copied from Python 3
  2. NamedFileInTemporaryDirectory and TemporaryWorkingDirectory from IPython, which
  3. uses the 3-clause BSD license.
  4. """
  5. from __future__ import print_function
  6. import os as _os
  7. import warnings as _warnings
  8. import sys as _sys
  9. # This code should only be used in Python versions < 3.2, since after that we
  10. # can rely on the stdlib itself.
  11. try:
  12. from tempfile import TemporaryDirectory
  13. except ImportError:
  14. from tempfile import mkdtemp, template
  15. class TemporaryDirectory(object):
  16. """Create and return a temporary directory. This has the same
  17. behavior as mkdtemp but can be used as a context manager. For
  18. example:
  19. with TemporaryDirectory() as tmpdir:
  20. ...
  21. Upon exiting the context, the directory and everything contained
  22. in it are removed.
  23. """
  24. def __init__(self, suffix="", prefix=template, dir=None):
  25. self.name = mkdtemp(suffix, prefix, dir)
  26. self._closed = False
  27. def __enter__(self):
  28. return self.name
  29. def cleanup(self, _warn=False):
  30. if self.name and not self._closed:
  31. try:
  32. self._rmtree(self.name)
  33. except (TypeError, AttributeError) as ex:
  34. # Issue #10188: Emit a warning on stderr
  35. # if the directory could not be cleaned
  36. # up due to missing globals
  37. if "None" not in str(ex):
  38. raise
  39. print("ERROR: {!r} while cleaning up {!r}".format(ex, self,),
  40. file=_sys.stderr)
  41. return
  42. self._closed = True
  43. if _warn:
  44. self._warn("Implicitly cleaning up {!r}".format(self),
  45. Warning)
  46. def __exit__(self, exc, value, tb):
  47. self.cleanup()
  48. def __del__(self):
  49. # Issue a ResourceWarning if implicit cleanup needed
  50. self.cleanup(_warn=True)
  51. # XXX (ncoghlan): The following code attempts to make
  52. # this class tolerant of the module nulling out process
  53. # that happens during CPython interpreter shutdown
  54. # Alas, it doesn't actually manage it. See issue #10188
  55. _listdir = staticmethod(_os.listdir)
  56. _path_join = staticmethod(_os.path.join)
  57. _isdir = staticmethod(_os.path.isdir)
  58. _remove = staticmethod(_os.remove)
  59. _rmdir = staticmethod(_os.rmdir)
  60. _os_error = _os.error
  61. _warn = _warnings.warn
  62. def _rmtree(self, path):
  63. # Essentially a stripped down version of shutil.rmtree. We can't
  64. # use globals because they may be None'ed out at shutdown.
  65. for name in self._listdir(path):
  66. fullname = self._path_join(path, name)
  67. try:
  68. isdir = self._isdir(fullname)
  69. except self._os_error:
  70. isdir = False
  71. if isdir:
  72. self._rmtree(fullname)
  73. else:
  74. try:
  75. self._remove(fullname)
  76. except self._os_error:
  77. pass
  78. try:
  79. self._rmdir(path)
  80. except self._os_error:
  81. pass
  82. class NamedFileInTemporaryDirectory(object):
  83. """Open a file named `filename` in a temporary directory.
  84. This context manager is preferred over :class:`tempfile.NamedTemporaryFile`
  85. when one needs to reopen the file, because on Windows only one handle on a
  86. file can be open at a time. You can close the returned handle explicitly
  87. inside the context without deleting the file, and the context manager will
  88. delete the whole directory when it exits.
  89. Arguments `mode` and `bufsize` are passed to `open`.
  90. Rest of the arguments are passed to `TemporaryDirectory`.
  91. Usage example::
  92. with NamedFileInTemporaryDirectory('myfile', 'wb') as f:
  93. f.write('stuff')
  94. f.close()
  95. # You can now pass f.name to things that will re-open the file
  96. """
  97. def __init__(self, filename, mode='w+b', bufsize=-1, **kwds):
  98. self._tmpdir = TemporaryDirectory(**kwds)
  99. path = _os.path.join(self._tmpdir.name, filename)
  100. self.file = open(path, mode, bufsize)
  101. def cleanup(self):
  102. self.file.close()
  103. self._tmpdir.cleanup()
  104. __del__ = cleanup
  105. def __enter__(self):
  106. return self.file
  107. def __exit__(self, type, value, traceback):
  108. self.cleanup()
  109. class TemporaryWorkingDirectory(TemporaryDirectory):
  110. """
  111. Creates a temporary directory and sets the cwd to that directory.
  112. Automatically reverts to previous cwd upon cleanup.
  113. Usage example::
  114. with TemporaryWorkingDirectory() as tmpdir:
  115. ...
  116. """
  117. def __enter__(self):
  118. self.old_wd = _os.getcwd()
  119. _os.chdir(self.name)
  120. return super(TemporaryWorkingDirectory, self).__enter__()
  121. def __exit__(self, exc, value, tb):
  122. _os.chdir(self.old_wd)
  123. return super(TemporaryWorkingDirectory, self).__exit__(exc, value, tb)