self_outdated_check.py 8.3 KB

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