tempdir.py 4.7 KB

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