req_tracker.py 4.3 KB

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