temp_dir.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172
  1. # The following comment should be removed at some point in the future.
  2. # mypy: disallow-untyped-defs=False
  3. from __future__ import absolute_import
  4. import errno
  5. import itertools
  6. import logging
  7. import os.path
  8. import tempfile
  9. from pip._internal.utils.misc import rmtree
  10. from pip._internal.utils.typing import MYPY_CHECK_RUNNING
  11. if MYPY_CHECK_RUNNING:
  12. from typing import Optional
  13. logger = logging.getLogger(__name__)
  14. class TempDirectory(object):
  15. """Helper class that owns and cleans up a temporary directory.
  16. This class can be used as a context manager or as an OO representation of a
  17. temporary directory.
  18. Attributes:
  19. path
  20. Location to the created temporary directory
  21. delete
  22. Whether the directory should be deleted when exiting
  23. (when used as a contextmanager)
  24. Methods:
  25. cleanup()
  26. Deletes the temporary directory
  27. When used as a context manager, if the delete attribute is True, on
  28. exiting the context the temporary directory is deleted.
  29. """
  30. def __init__(
  31. self,
  32. path=None, # type: Optional[str]
  33. delete=None, # type: Optional[bool]
  34. kind="temp"
  35. ):
  36. super(TempDirectory, self).__init__()
  37. if path is None and delete is None:
  38. # If we were not given an explicit directory, and we were not given
  39. # an explicit delete option, then we'll default to deleting.
  40. delete = True
  41. if path is None:
  42. path = self._create(kind)
  43. self._path = path
  44. self._deleted = False
  45. self.delete = delete
  46. self.kind = kind
  47. @property
  48. def path(self):
  49. # type: () -> str
  50. assert not self._deleted, (
  51. "Attempted to access deleted path: {}".format(self._path)
  52. )
  53. return self._path
  54. def __repr__(self):
  55. return "<{} {!r}>".format(self.__class__.__name__, self.path)
  56. def __enter__(self):
  57. return self
  58. def __exit__(self, exc, value, tb):
  59. if self.delete:
  60. self.cleanup()
  61. def _create(self, kind):
  62. """Create a temporary directory and store its path in self.path
  63. """
  64. # We realpath here because some systems have their default tmpdir
  65. # symlinked to another directory. This tends to confuse build
  66. # scripts, so we canonicalize the path by traversing potential
  67. # symlinks here.
  68. path = os.path.realpath(
  69. tempfile.mkdtemp(prefix="pip-{}-".format(kind))
  70. )
  71. logger.debug("Created temporary directory: {}".format(path))
  72. return path
  73. def cleanup(self):
  74. """Remove the temporary directory created and reset state
  75. """
  76. self._deleted = True
  77. if os.path.exists(self._path):
  78. rmtree(self._path)
  79. class AdjacentTempDirectory(TempDirectory):
  80. """Helper class that creates a temporary directory adjacent to a real one.
  81. Attributes:
  82. original
  83. The original directory to create a temp directory for.
  84. path
  85. After calling create() or entering, contains the full
  86. path to the temporary directory.
  87. delete
  88. Whether the directory should be deleted when exiting
  89. (when used as a contextmanager)
  90. """
  91. # The characters that may be used to name the temp directory
  92. # We always prepend a ~ and then rotate through these until
  93. # a usable name is found.
  94. # pkg_resources raises a different error for .dist-info folder
  95. # with leading '-' and invalid metadata
  96. LEADING_CHARS = "-~.=%0123456789"
  97. def __init__(self, original, delete=None):
  98. self.original = original.rstrip('/\\')
  99. super(AdjacentTempDirectory, self).__init__(delete=delete)
  100. @classmethod
  101. def _generate_names(cls, name):
  102. """Generates a series of temporary names.
  103. The algorithm replaces the leading characters in the name
  104. with ones that are valid filesystem characters, but are not
  105. valid package names (for both Python and pip definitions of
  106. package).
  107. """
  108. for i in range(1, len(name)):
  109. for candidate in itertools.combinations_with_replacement(
  110. cls.LEADING_CHARS, i - 1):
  111. new_name = '~' + ''.join(candidate) + name[i:]
  112. if new_name != name:
  113. yield new_name
  114. # If we make it this far, we will have to make a longer name
  115. for i in range(len(cls.LEADING_CHARS)):
  116. for candidate in itertools.combinations_with_replacement(
  117. cls.LEADING_CHARS, i):
  118. new_name = '~' + ''.join(candidate) + name
  119. if new_name != name:
  120. yield new_name
  121. def _create(self, kind):
  122. root, name = os.path.split(self.original)
  123. for candidate in self._generate_names(name):
  124. path = os.path.join(root, candidate)
  125. try:
  126. os.mkdir(path)
  127. except OSError as ex:
  128. # Continue if the name exists already
  129. if ex.errno != errno.EEXIST:
  130. raise
  131. else:
  132. path = os.path.realpath(path)
  133. break
  134. else:
  135. # Final fallback on the default behavior.
  136. path = os.path.realpath(
  137. tempfile.mkdtemp(prefix="pip-{}-".format(kind))
  138. )
  139. logger.debug("Created temporary directory: {}".format(path))
  140. return path