| 1 | # Copyright 2026 Google LLC |
| 2 | # |
| 3 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | # you may not use this file except in compliance with the License. |
| 5 | # You may obtain a copy of the License at |
| 6 | # |
| 7 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | # |
| 9 | # Unless required by applicable law or agreed to in writing, software |
| 10 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | # See the License for the specific language governing permissions and |
| 13 | # limitations under the License. |
| 14 | |
| 15 | """Auto-update subsystem. |
| 16 | |
| 17 | Owns version detection, the PyPI-style update probe, the on-disk |
| 18 | ``latest_version`` cache, and the upgrade-banner UX. The CLI's global |
| 19 | callback (``cli.py``) calls ``check_for_updates`` once per day and |
| 20 | ``maybe_show_cached_banner`` on every other invocation; the |
| 21 | ``colab update`` Typer command (``commands/utility.py``) delegates to |
| 22 | ``check_for_updates``. |
| 23 | """ |
| 24 | |
| 25 | import json |
| 26 | import platform |
| 27 | import subprocess |
| 28 | import urllib.request |
| 29 | from datetime import datetime, timezone |
| 30 | from importlib.metadata import PackageNotFoundError, version as installed_version |
| 31 | from packaging.version import InvalidVersion, Version |
| 32 | from typing import Optional |
| 33 | |
| 34 | import typer |
| 35 | |
| 36 | from colab_cli.common import state |
| 37 | from colab_cli.state import Settings |
| 38 | |
| 39 | # PyPI distribution name (different from the importable package name `colab`). |
| 40 | PYPI_PACKAGE_NAME = "google-colab-cli" |
| 41 | |
| 42 | |
| 43 | # ---------- Version detection ------------------------------------------- |
| 44 | |
| 45 | |
| 46 | def is_self_install_supported() -> bool: |
| 47 | """Return True if self-install (--install) is supported on the current platform.""" |
| 48 | return platform.system() in ("Linux", "Darwin") |
| 49 | |
| 50 | |
| 51 | def get_app_version() -> str: |
| 52 | """Return the installed package version, falling back to the git short hash.""" |
| 53 | try: |
| 54 | return installed_version("google-colab-cli") |
| 55 | except (PackageNotFoundError, InvalidVersion): |
| 56 | pass |
| 57 | |
| 58 | try: |
| 59 | return subprocess.check_output( |
| 60 | ["git", "rev-parse", "--short", "HEAD"], |
| 61 | stderr=subprocess.DEVNULL, |
| 62 | encoding="utf-8", |
| 63 | ).strip() |
| 64 | except Exception: |
| 65 | return "unknown" |
| 66 | |
| 67 | |
| 68 | # ---------- Source fetchers --------------------------------------------- |
| 69 | |
| 70 | |
| 71 | def _parse_version(payload: Optional[dict]) -> Optional[str]: |
| 72 | """Returns ``info.version`` from a PyPI-style payload, or None.""" |
| 73 | return (payload or {}).get("info", {}).get("version") |
| 74 | |
| 75 | |
| 76 | def _fetch_pypi(url: str, quiet: bool) -> Optional[dict]: |
| 77 | """Fetches and parses the PyPI-style JSON document at ``url``.""" |
| 78 | try: |
| 79 | with urllib.request.urlopen(url, timeout=5) as response: |
| 80 | return json.loads(response.read().decode("utf-8")) |
| 81 | except Exception as e: |
| 82 | if not quiet: |
| 83 | typer.echo(f"[colab] Warning: Failed to fetch update info: {e}") |
| 84 | return None |
| 85 | |
| 86 | |
| 87 | # ---------- Version comparison ------------------------------------------ |
| 88 | |
| 89 | |
| 90 | def _is_newer(candidate: Optional[str], current: str) -> bool: |
| 91 | """True when ``candidate`` strictly exceeds ``current`` (PEP 440).""" |
| 92 | if not candidate: |
| 93 | return False |
| 94 | try: |
| 95 | return Version(candidate) > Version(current) |
| 96 | except InvalidVersion: |
| 97 | return candidate != current |
| 98 | |
| 99 | |
| 100 | # ---------- UX ---------------------------------------------------------- |
| 101 | |
| 102 | |
| 103 | def announce_upgrade( |
| 104 | latest: str, |
| 105 | current: str, |
| 106 | install_cmd: str, |
| 107 | *, |
| 108 | show_disable_hint: bool = False, |
| 109 | ) -> None: |
| 110 | """Print the upgrade banner. |
| 111 | |
| 112 | ``show_disable_hint`` controls whether the trailing line that explains |
| 113 | how to silence the auto-check is included. It is only added when the |
| 114 | banner is shown unsolicited (the daily background fetch and the cached |
| 115 | banner on subsequent invocations); explicit ``colab update`` calls |
| 116 | omit it because the user already opted in to seeing the result. |
| 117 | """ |
| 118 | typer.echo( |
| 119 | f"\n[colab] A new version of Colab CLI is available: {latest} (current: {current})" |
| 120 | ) |
| 121 | if is_self_install_supported() and ("pip" in install_cmd or "uv" in install_cmd): |
| 122 | typer.echo("[colab] You can run 'colab update --install' to upgrade in place.") |
| 123 | typer.echo(f"[colab] Run '{install_cmd}' to update.") |
| 124 | if show_disable_hint: |
| 125 | typer.echo( |
| 126 | "[colab] To silence this check, set " |
| 127 | '"enable_update_check": false in ' |
| 128 | "~/.config/colab-cli/settings.json" |
| 129 | ) |
| 130 | typer.echo("") |
| 131 | |
| 132 | |
| 133 | # ---------- Orchestration ----------------------------------------------- |
| 134 | |
| 135 | |
| 136 | def _get_install_command() -> str: |
| 137 | """Return the recommended installation command based on the environment.""" |
| 138 | import sys |
| 139 | |
| 140 | if is_self_install_supported() and "/uv/tools/" in sys.executable: |
| 141 | return f"uv tool install -U {PYPI_PACKAGE_NAME}" |
| 142 | return f"pip install --upgrade {PYPI_PACKAGE_NAME}" |
| 143 | |
| 144 | |
| 145 | def check_for_updates(quiet: bool = False) -> None: |
| 146 | """Check PyPI for updates and print a message if a new version is available. |
| 147 | |
| 148 | The disable-hint is appended to the banner only when ``quiet`` is True |
| 149 | (the daily background fetch); explicit ``colab update`` invocations |
| 150 | (``quiet=False``) omit it because the user asked for the check. |
| 151 | """ |
| 152 | settings = state.settings_store.load() |
| 153 | current = get_app_version() |
| 154 | |
| 155 | try: |
| 156 | pypi = _fetch_pypi(settings.update_url, quiet) |
| 157 | pypi_v = _parse_version(pypi) |
| 158 | |
| 159 | if _is_newer(pypi_v, current): |
| 160 | announce_upgrade( |
| 161 | pypi_v, |
| 162 | current, |
| 163 | _get_install_command(), |
| 164 | show_disable_hint=quiet, |
| 165 | ) |
| 166 | elif not quiet: |
| 167 | suffix = f", latest: {pypi_v}" if pypi_v else "" |
| 168 | typer.echo(f"[colab] Colab CLI is up to date (version: {current}{suffix}).") |
| 169 | |
| 170 | # Cache the highest observed version; never downgrade. |
| 171 | cached = settings.latest_version or "0" |
| 172 | if _is_newer(pypi_v, cached): |
| 173 | settings.latest_version = pypi_v |
| 174 | |
| 175 | settings.last_check = datetime.now(timezone.utc) |
| 176 | state.settings_store.save(settings) |
| 177 | |
| 178 | except Exception as e: |
| 179 | if not quiet: |
| 180 | typer.echo(f"[colab] Failed to check for updates: {e}") |
| 181 | |
| 182 | |
| 183 | # ---------- Background hooks (called from cli.py) ----------------------- |
| 184 | |
| 185 | |
| 186 | def _is_throttled(settings: Settings, *, now: Optional[datetime] = None) -> bool: |
| 187 | """True if the once-per-day fetch should be skipped.""" |
| 188 | if settings.last_check is None: |
| 189 | return False |
| 190 | now = now or datetime.now(timezone.utc) |
| 191 | return (now - settings.last_check).days < 1 |
| 192 | |
| 193 | |
| 194 | def maybe_show_cached_banner(settings: Settings) -> None: |
| 195 | """Print the cached upgrade banner if the cache reports a newer version. |
| 196 | |
| 197 | Called from the global CLI callback when the daily fetch is throttled. |
| 198 | The banner uses a generic ``colab update`` install hint because the |
| 199 | cache does not record which source supplied the version; the disable |
| 200 | hint is shown because this is unsolicited output. |
| 201 | """ |
| 202 | if not settings.latest_version: |
| 203 | return |
| 204 | current = get_app_version() |
| 205 | if not _is_newer(settings.latest_version, current): |
| 206 | return |
| 207 | announce_upgrade( |
| 208 | settings.latest_version, |
| 209 | current, |
| 210 | "colab update", |
| 211 | show_disable_hint=True, |
| 212 | ) |
| 213 | |
| 214 | |
| 215 | def run_background_check() -> None: |
| 216 | """Entry point for the global CLI callback. |
| 217 | |
| 218 | Performs either the daily fetch (which writes the cache) or, if |
| 219 | throttled, surfaces the cached banner. Honors the |
| 220 | ``enable_update_check`` master switch. |
| 221 | """ |
| 222 | settings = state.settings_store.load() |
| 223 | if not settings.enable_update_check: |
| 224 | return |
| 225 | if _is_throttled(settings): |
| 226 | maybe_show_cached_banner(settings) |
| 227 | else: |
| 228 | check_for_updates(quiet=True) |
| 229 | |
| 230 | |
| 231 | # ---------- Self-install ------------------------------------------------ |
| 232 | |
| 233 | |
| 234 | def self_install() -> None: |
| 235 | """Upgrade the CLI in place, detecting uv vs pip.""" |
| 236 | import sys |
| 237 | |
| 238 | # If the executable path contains "/uv/", we assume it was installed via |
| 239 | # `uv tool install` and use `uv` to upgrade it. |
| 240 | if "/uv/tools/" in sys.executable: |
| 241 | cmd = ["uv", "tool", "install", "-U", PYPI_PACKAGE_NAME] |
| 242 | else: |
| 243 | cmd = [sys.executable, "-m", "pip", "install", "-U", PYPI_PACKAGE_NAME] |
| 244 | |
| 245 | typer.echo(f"[colab] Running: {' '.join(cmd)}") |
| 246 | result = subprocess.run(cmd) |
| 247 | if result.returncode != 0: |
| 248 | raise typer.Exit(code=result.returncode) |