self_outdated_check.py 8.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255
  1. from __future__ import annotations
  2. import datetime
  3. import functools
  4. import hashlib
  5. import json
  6. import logging
  7. import optparse
  8. import os.path
  9. import sys
  10. from dataclasses import dataclass
  11. from typing import Callable
  12. from pip._vendor.packaging.version import Version
  13. from pip._vendor.packaging.version import parse as parse_version
  14. from pip._vendor.rich.console import Group
  15. from pip._vendor.rich.markup import escape
  16. from pip._vendor.rich.text import Text
  17. from pip._internal.index.collector import LinkCollector
  18. from pip._internal.index.package_finder import PackageFinder
  19. from pip._internal.metadata import get_default_environment
  20. from pip._internal.models.release_control import ReleaseControl
  21. from pip._internal.models.selection_prefs import SelectionPreferences
  22. from pip._internal.network.session import PipSession
  23. from pip._internal.utils.compat import WINDOWS
  24. from pip._internal.utils.datetime import parse_iso_datetime
  25. from pip._internal.utils.entrypoints import (
  26. get_best_invocation_for_this_pip,
  27. get_best_invocation_for_this_python,
  28. )
  29. from pip._internal.utils.filesystem import (
  30. adjacent_tmp_file,
  31. check_path_owner,
  32. copy_directory_permissions,
  33. replace,
  34. )
  35. from pip._internal.utils.misc import (
  36. ExternallyManagedEnvironment,
  37. check_externally_managed,
  38. ensure_dir,
  39. )
  40. _WEEK = datetime.timedelta(days=7)
  41. logger = logging.getLogger(__name__)
  42. def _get_statefile_name(key: str) -> str:
  43. key_bytes = key.encode()
  44. name = hashlib.sha224(key_bytes).hexdigest()
  45. return name
  46. class SelfCheckState:
  47. def __init__(self, cache_dir: str) -> None:
  48. self._state: dict[str, str] = {}
  49. self._statefile_path = None
  50. # Try to load the existing state
  51. if cache_dir:
  52. self._statefile_path = os.path.join(
  53. cache_dir, "selfcheck", _get_statefile_name(self.key)
  54. )
  55. try:
  56. with open(self._statefile_path, encoding="utf-8") as statefile:
  57. self._state = json.load(statefile)
  58. except (OSError, ValueError, KeyError):
  59. # Explicitly suppressing exceptions, since we don't want to
  60. # error out if the cache file is invalid.
  61. pass
  62. @property
  63. def key(self) -> str:
  64. return sys.prefix
  65. def get(self, current_time: datetime.datetime) -> str | None:
  66. """Check if we have a not-outdated version loaded already."""
  67. if not self._state:
  68. return None
  69. if "last_check" not in self._state:
  70. return None
  71. if "pypi_version" not in self._state:
  72. return None
  73. # Determine if we need to refresh the state
  74. last_check = parse_iso_datetime(self._state["last_check"])
  75. time_since_last_check = current_time - last_check
  76. if time_since_last_check > _WEEK:
  77. return None
  78. return self._state["pypi_version"]
  79. def set(self, pypi_version: str, current_time: datetime.datetime) -> None:
  80. # If we do not have a path to cache in, don't bother saving.
  81. if not self._statefile_path:
  82. return
  83. statefile_directory = os.path.dirname(self._statefile_path)
  84. # Check to make sure that we own the directory
  85. if not check_path_owner(statefile_directory):
  86. return
  87. # Now that we've ensured the directory is owned by this user, we'll go
  88. # ahead and make sure that all our directories are created.
  89. ensure_dir(statefile_directory)
  90. state = {
  91. # Include the key so it's easy to tell which pip wrote the
  92. # file.
  93. "key": self.key,
  94. "last_check": current_time.isoformat(),
  95. "pypi_version": pypi_version,
  96. }
  97. text = json.dumps(state, sort_keys=True, separators=(",", ":"))
  98. with adjacent_tmp_file(self._statefile_path) as f:
  99. f.write(text.encode())
  100. copy_directory_permissions(statefile_directory, f)
  101. try:
  102. # Since we have a prefix-specific state file, we can just
  103. # overwrite whatever is there, no need to check.
  104. replace(f.name, self._statefile_path)
  105. except OSError:
  106. # Best effort.
  107. pass
  108. @dataclass
  109. class UpgradePrompt:
  110. old: str
  111. new: str
  112. def __rich__(self) -> Group:
  113. if WINDOWS:
  114. pip_cmd = f"{get_best_invocation_for_this_python()} -m pip"
  115. else:
  116. pip_cmd = get_best_invocation_for_this_pip()
  117. notice = "[bold][[reset][blue]notice[reset][bold]][reset]"
  118. return Group(
  119. Text(),
  120. Text.from_markup(
  121. f"{notice} A new release of pip is available: "
  122. f"[red]{self.old}[reset] -> [green]{self.new}[reset]"
  123. ),
  124. Text.from_markup(
  125. f"{notice} To update, run: "
  126. f"[green]{escape(pip_cmd)} install --upgrade pip"
  127. ),
  128. )
  129. def was_installed_by_pip(pkg: str) -> bool:
  130. """Checks whether pkg was installed by pip
  131. This is used not to display the upgrade message when pip is in fact
  132. installed by system package manager, such as dnf on Fedora.
  133. """
  134. dist = get_default_environment().get_distribution(pkg)
  135. return dist is not None and "pip" == dist.installer
  136. def _get_current_remote_pip_version(
  137. session: PipSession, options: optparse.Values
  138. ) -> str | None:
  139. # Lets use PackageFinder to see what the latest pip version is
  140. link_collector = LinkCollector.create(
  141. session,
  142. options=options,
  143. suppress_no_index=True,
  144. )
  145. # Pass allow_yanked=False so we don't suggest upgrading to a
  146. # yanked version.
  147. selection_prefs = SelectionPreferences(
  148. allow_yanked=False,
  149. release_control=ReleaseControl(only_final={"pip"}),
  150. )
  151. finder = PackageFinder.create(
  152. link_collector=link_collector,
  153. selection_prefs=selection_prefs,
  154. )
  155. best_candidate = finder.find_best_candidate("pip").best_candidate
  156. if best_candidate is None:
  157. return None
  158. return str(best_candidate.version)
  159. def _self_version_check_logic(
  160. *,
  161. state: SelfCheckState,
  162. current_time: datetime.datetime,
  163. local_version: Version,
  164. get_remote_version: Callable[[], str | None],
  165. ) -> UpgradePrompt | None:
  166. remote_version_str = state.get(current_time)
  167. if remote_version_str is None:
  168. remote_version_str = get_remote_version()
  169. if remote_version_str is None:
  170. logger.debug("No remote pip version found")
  171. return None
  172. state.set(remote_version_str, current_time)
  173. remote_version = parse_version(remote_version_str)
  174. logger.debug("Remote version of pip: %s", remote_version)
  175. logger.debug("Local version of pip: %s", local_version)
  176. pip_installed_by_pip = was_installed_by_pip("pip")
  177. logger.debug("Was pip installed by pip? %s", pip_installed_by_pip)
  178. if not pip_installed_by_pip:
  179. return None # Only suggest upgrade if pip is installed by pip.
  180. local_version_is_older = (
  181. local_version < remote_version
  182. and local_version.base_version != remote_version.base_version
  183. )
  184. if local_version_is_older:
  185. return UpgradePrompt(old=str(local_version), new=remote_version_str)
  186. return None
  187. def pip_self_version_check(session: PipSession, options: optparse.Values) -> None:
  188. """Check for an update for pip.
  189. Limit the frequency of checks to once per week. State is stored either in
  190. the active virtualenv or in the user's USER_CACHE_DIR keyed off the prefix
  191. of the pip script path.
  192. """
  193. installed_dist = get_default_environment().get_distribution("pip")
  194. if not installed_dist:
  195. return
  196. try:
  197. check_externally_managed()
  198. except ExternallyManagedEnvironment:
  199. return
  200. upgrade_prompt = _self_version_check_logic(
  201. state=SelfCheckState(cache_dir=options.cache_dir),
  202. current_time=datetime.datetime.now(datetime.timezone.utc),
  203. local_version=installed_dist.version,
  204. get_remote_version=functools.partial(
  205. _get_current_remote_pip_version, session, options
  206. ),
  207. )
  208. if upgrade_prompt is not None:
  209. logger.warning("%s", upgrade_prompt, extra={"rich": True})