req_tracker.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  1. # The following comment should be removed at some point in the future.
  2. # mypy: strict-optional=False
  3. from __future__ import absolute_import
  4. import contextlib
  5. import errno
  6. import hashlib
  7. import logging
  8. import os
  9. from pip._vendor import contextlib2
  10. from pip._internal.utils.temp_dir import TempDirectory
  11. from pip._internal.utils.typing import MYPY_CHECK_RUNNING
  12. if MYPY_CHECK_RUNNING:
  13. from types import TracebackType
  14. from typing import Dict, Iterator, Optional, Set, Type, Union
  15. from pip._internal.req.req_install import InstallRequirement
  16. from pip._internal.models.link import Link
  17. logger = logging.getLogger(__name__)
  18. @contextlib.contextmanager
  19. def update_env_context_manager(**changes):
  20. # type: (str) -> Iterator[None]
  21. target = os.environ
  22. # Save values from the target and change them.
  23. non_existent_marker = object()
  24. saved_values = {} # type: Dict[str, Union[object, str]]
  25. for name, new_value in changes.items():
  26. try:
  27. saved_values[name] = target[name]
  28. except KeyError:
  29. saved_values[name] = non_existent_marker
  30. target[name] = new_value
  31. try:
  32. yield
  33. finally:
  34. # Restore original values in the target.
  35. for name, original_value in saved_values.items():
  36. if original_value is non_existent_marker:
  37. del target[name]
  38. else:
  39. assert isinstance(original_value, str) # for mypy
  40. target[name] = original_value
  41. @contextlib.contextmanager
  42. def get_requirement_tracker():
  43. # type: () -> Iterator[RequirementTracker]
  44. root = os.environ.get('PIP_REQ_TRACKER')
  45. with contextlib2.ExitStack() as ctx:
  46. if root is None:
  47. root = ctx.enter_context(
  48. TempDirectory(kind='req-tracker')
  49. ).path
  50. ctx.enter_context(update_env_context_manager(PIP_REQ_TRACKER=root))
  51. logger.debug("Initialized build tracking at %s", root)
  52. with RequirementTracker(root) as tracker:
  53. yield tracker
  54. class RequirementTracker(object):
  55. def __init__(self, root):
  56. # type: (str) -> None
  57. self._root = root
  58. self._entries = set() # type: Set[InstallRequirement]
  59. logger.debug("Created build tracker: %s", self._root)
  60. def __enter__(self):
  61. # type: () -> RequirementTracker
  62. logger.debug("Entered build tracker: %s", self._root)
  63. return self
  64. def __exit__(
  65. self,
  66. exc_type, # type: Optional[Type[BaseException]]
  67. exc_val, # type: Optional[BaseException]
  68. exc_tb # type: Optional[TracebackType]
  69. ):
  70. # type: (...) -> None
  71. self.cleanup()
  72. def _entry_path(self, link):
  73. # type: (Link) -> str
  74. hashed = hashlib.sha224(link.url_without_fragment.encode()).hexdigest()
  75. return os.path.join(self._root, hashed)
  76. def add(self, req):
  77. # type: (InstallRequirement) -> None
  78. """Add an InstallRequirement to build tracking.
  79. """
  80. # Get the file to write information about this requirement.
  81. entry_path = self._entry_path(req.link)
  82. # Try reading from the file. If it exists and can be read from, a build
  83. # is already in progress, so a LookupError is raised.
  84. try:
  85. with open(entry_path) as fp:
  86. contents = fp.read()
  87. except IOError as e:
  88. # if the error is anything other than "file does not exist", raise.
  89. if e.errno != errno.ENOENT:
  90. raise
  91. else:
  92. message = '{} is already being built: {}'.format(
  93. req.link, contents)
  94. raise LookupError(message)
  95. # If we're here, req should really not be building already.
  96. assert req not in self._entries
  97. # Start tracking this requirement.
  98. with open(entry_path, 'w') as fp:
  99. fp.write(str(req))
  100. self._entries.add(req)
  101. logger.debug('Added %s to build tracker %r', req, self._root)
  102. def remove(self, req):
  103. # type: (InstallRequirement) -> None
  104. """Remove an InstallRequirement from build tracking.
  105. """
  106. # Delete the created file and the corresponding entries.
  107. os.unlink(self._entry_path(req.link))
  108. self._entries.remove(req)
  109. logger.debug('Removed %s from build tracker %r', req, self._root)
  110. def cleanup(self):
  111. # type: () -> None
  112. for req in set(self._entries):
  113. self.remove(req)
  114. logger.debug("Removed build tracker: %r", self._root)
  115. @contextlib.contextmanager
  116. def track(self, req):
  117. # type: (InstallRequirement) -> Iterator[None]
  118. self.add(req)
  119. yield
  120. self.remove(req)