125 lines
4.0 KiB
Python
125 lines
4.0 KiB
Python
|
import contextlib
|
||
|
|
import hashlib
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
from types import TracebackType
|
||
|
|
from typing import Dict, Iterator, Optional, Set, Type, Union
|
||
|
|
|
||
|
|
from pip._internal.models.link import Link
|
||
|
|
from pip._internal.req.req_install import InstallRequirement
|
||
|
|
from pip._internal.utils.temp_dir import TempDirectory
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
@contextlib.contextmanager
|
||
|
|
def update_env_context_manager(**changes: str) -> Iterator[None]:
|
||
|
|
target = os.environ
|
||
|
|
|
||
|
|
# Save values from the target and change them.
|
||
|
|
non_existent_marker = object()
|
||
|
|
saved_values: Dict[str, Union[object, str]] = {}
|
||
|
|
for name, new_value in changes.items():
|
||
|
|
try:
|
||
|
|
saved_values[name] = target[name]
|
||
|
|
except KeyError:
|
||
|
|
saved_values[name] = non_existent_marker
|
||
|
|
target[name] = new_value
|
||
|
|
|
||
|
|
try:
|
||
|
|
yield
|
||
|
|
finally:
|
||
|
|
# Restore original values in the target.
|
||
|
|
for name, original_value in saved_values.items():
|
||
|
|
if original_value is non_existent_marker:
|
||
|
|
del target[name]
|
||
|
|
else:
|
||
|
|
assert isinstance(original_value, str) # for mypy
|
||
|
|
target[name] = original_value
|
||
|
|
|
||
|
|
|
||
|
|
@contextlib.contextmanager
|
||
|
|
def get_requirement_tracker() -> Iterator["RequirementTracker"]:
|
||
|
|
root = os.environ.get("PIP_REQ_TRACKER")
|
||
|
|
with contextlib.ExitStack() as ctx:
|
||
|
|
if root is None:
|
||
|
|
root = ctx.enter_context(TempDirectory(kind="req-tracker")).path
|
||
|
|
ctx.enter_context(update_env_context_manager(PIP_REQ_TRACKER=root))
|
||
|
|
logger.debug("Initialized build tracking at %s", root)
|
||
|
|
|
||
|
|
with RequirementTracker(root) as tracker:
|
||
|
|
yield tracker
|
||
|
|
|
||
|
|
|
||
|
|
class RequirementTracker:
|
||
|
|
def __init__(self, root: str) -> None:
|
||
|
|
self._root = root
|
||
|
|
self._entries: Set[InstallRequirement] = set()
|
||
|
|
logger.debug("Created build tracker: %s", self._root)
|
||
|
|
|
||
|
|
def __enter__(self) -> "RequirementTracker":
|
||
|
|
logger.debug("Entered build tracker: %s", self._root)
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(
|
||
|
|
self,
|
||
|
|
exc_type: Optional[Type[BaseException]],
|
||
|
|
exc_val: Optional[BaseException],
|
||
|
|
exc_tb: Optional[TracebackType],
|
||
|
|
) -> None:
|
||
|
|
self.cleanup()
|
||
|
|
|
||
|
|
def _entry_path(self, link: Link) -> str:
|
||
|
|
hashed = hashlib.sha224(link.url_without_fragment.encode()).hexdigest()
|
||
|
|
return os.path.join(self._root, hashed)
|
||
|
|
|
||
|
|
def add(self, req: InstallRequirement) -> None:
|
||
|
|
"""Add an InstallRequirement to build tracking."""
|
||
|
|
|
||
|
|
assert req.link
|
||
|
|
# Get the file to write information about this requirement.
|
||
|
|
entry_path = self._entry_path(req.link)
|
||
|
|
|
||
|
|
# Try reading from the file. If it exists and can be read from, a build
|
||
|
|
# is already in progress, so a LookupError is raised.
|
||
|
|
try:
|
||
|
|
with open(entry_path) as fp:
|
||
|
|
contents = fp.read()
|
||
|
|
except FileNotFoundError:
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
message = "{} is already being built: {}".format(req.link, contents)
|
||
|
|
raise LookupError(message)
|
||
|
|
|
||
|
|
# If we're here, req should really not be building already.
|
||
|
|
assert req not in self._entries
|
||
|
|
|
||
|
|
# Start tracking this requirement.
|
||
|
|
with open(entry_path, "w", encoding="utf-8") as fp:
|
||
|
|
fp.write(str(req))
|
||
|
|
self._entries.add(req)
|
||
|
|
|
||
|
|
logger.debug("Added %s to build tracker %r", req, self._root)
|
||
|
|
|
||
|
|
def remove(self, req: InstallRequirement) -> None:
|
||
|
|
"""Remove an InstallRequirement from build tracking."""
|
||
|
|
|
||
|
|
assert req.link
|
||
|
|
# Delete the created file and the corresponding entries.
|
||
|
|
os.unlink(self._entry_path(req.link))
|
||
|
|
self._entries.remove(req)
|
||
|
|
|
||
|
|
logger.debug("Removed %s from build tracker %r", req, self._root)
|
||
|
|
|
||
|
|
def cleanup(self) -> None:
|
||
|
|
for req in set(self._entries):
|
||
|
|
self.remove(req)
|
||
|
|
|
||
|
|
logger.debug("Removed build tracker: %r", self._root)
|
||
|
|
|
||
|
|
@contextlib.contextmanager
|
||
|
|
def track(self, req: InstallRequirement) -> Iterator[None]:
|
||
|
|
self.add(req)
|
||
|
|
yield
|
||
|
|
self.remove(req)
|