self_outdated_check.py 8.1 KB

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