"""`headroom update` — self-update across the supported install methods. Detects how Headroom was installed and runs the matching upgrade command, or refuses with clear guidance when in-tool self-update isn't appropriate (git checkout, editable install, Docker image, system package manager). The upgrade always runs through ``sys.executable -m pip`` for the pip path so it can never touch a different interpreter than the one actually running Headroom. """ from __future__ import annotations import json import os import shlex import shutil import subprocess import sys from dataclasses import dataclass from pathlib import Path import click from headroom._subprocess import run as _subprocess_run from headroom.update_check import ( PACKAGE_NAME, fetch_latest_version, installed_version, write_cache, ) from .main import main @dataclass(frozen=True) class InstallMethod: """How Headroom is installed and how (or whether) to upgrade it.""" kind: str can_self_update: bool argv: list[str] | None = None guidance: str | None = None def _is_source_checkout() -> bool: try: from headroom._version import _source_root return _source_root() is not None except Exception: return False def _is_editable_install() -> bool: """Detect a PEP 660 editable install via its ``direct_url.json``.""" try: from importlib.metadata import distribution raw = distribution(PACKAGE_NAME).read_text("direct_url.json") if not raw: return False data = json.loads(raw) dir_info = data.get("dir_info") return bool(isinstance(dir_info, dict) and dir_info.get("editable")) except Exception: return False def _in_docker() -> bool: try: from pathlib import Path return Path("/.dockerenv").exists() or bool( os.environ.get("HEADROOM_IN_DOCKER", "").strip() ) except Exception: return False def _in_virtualenv() -> bool: """True inside a venv/virtualenv or a conda environment.""" if getattr(sys, "prefix", "") != getattr(sys, "base_prefix", ""): return True # Conda envs often share prefix==base_prefix for `base`; pip -U is still safe # inside any conda env, so treat an active CONDA_PREFIX as an environment. return bool(os.environ.get("CONDA_PREFIX", "").strip()) def _norm(path: str | os.PathLike[str] | None) -> str: """Resolve + lowercase a path for cross-platform substring matching.""" if not path: return "" try: from pathlib import Path return str(Path(path).resolve()).replace("\\", "/").lower() except Exception: return str(path).replace("\\", "/").lower() def _package_location() -> str | None: """Normalized base directory the headroom-ai distribution is installed in.""" try: from importlib.metadata import distribution return _norm(str(distribution(PACKAGE_NAME).locate_file(""))) except Exception: return None def _user_site() -> str | None: try: import site return _norm(site.getusersitepackages()) except Exception: return None def _is_user_site_install(location: str | None) -> bool: user = _user_site() if not user or not location: return False # Path-segment containment so "/.../site" never matches "/.../site-packages". return location == user or location.startswith(user.rstrip("/") + "/") def _format_cmd(argv: list[str]) -> str: """Render argv as a copy-pasteable shell string (handles spaces in paths).""" if sys.platform.startswith("win"): return subprocess.list2cmdline(argv) return shlex.join(argv) def _is_externally_managed() -> bool: """Detect a PEP 668 ``EXTERNALLY-MANAGED`` marker (Homebrew, Debian, etc.).""" try: import sysconfig from pathlib import Path for key in ("stdlib", "platstdlib", "purelib"): base = sysconfig.get_path(key) if base and Path(base, "EXTERNALLY-MANAGED").exists(): return True except Exception: return False return False def _spec(extras: str | None) -> str: extras = (extras or "").strip().strip("[]") return f"{PACKAGE_NAME}[{extras}]" if extras else PACKAGE_NAME def _managed_env_guidance() -> str: if sys.platform == "darwin": hint = "`brew upgrade headroom-ai` (if installed via Homebrew), or reinstall with pipx" elif sys.platform.startswith("win"): hint = "reinstall with pipx (`pipx install headroom-ai`) or use a virtualenv" else: hint = "use your distro package manager, or reinstall with pipx / a virtualenv" return ( "Headroom is installed in an externally-managed system Python (PEP 668). " f"Don't pip into it — {hint}." ) def _find_core_pyd() -> Path | None: """Locate the _core.pyd file inside the headroom site-packages directory.""" try: import headroom headroom_path = Path(headroom.__file__).parent core_pyd = headroom_path / "_core.pyd" return core_pyd if core_pyd.exists() else None except Exception: return None def _is_pyd_locked(pyd_path: Path) -> bool: """Verify whether the .pyd file is locked on Windows by attempting to open it in r+b mode.""" try: with open(pyd_path, "r+b") as f: f.close() return False except (PermissionError, OSError): return True def _make_backup(pyd_path: Path) -> Path: """Create a backup of the .pyd file as .pyd.bak.""" backup_path = pyd_path.with_suffix(pyd_path.suffix + ".bak") shutil.copy2(pyd_path, backup_path) return backup_path def _restore_backup(pyd_path: Path, backup_path: Path) -> None: """Restore the .pyd file atomically by replacing the original with the backup.""" os.replace(backup_path, pyd_path) def _test_core_integrity() -> bool: """Test if headroom._core can be imported successfully.""" try: result = _subprocess_run( [sys.executable, "-c", "from headroom._core import hello"], capture_output=True, text=True, timeout=10, ) return result.returncode == 0 except Exception: return False def safe_update(argv: list[str]) -> int: """Execute update with protection against .pyd corruption on Windows. Detects if _core.pyd is locked (indicating that headroom proxy is running), makes a backup, executes the update, tests integrity, and restores the backup in case of failure. Returns 0 if successful, error code otherwise. """ core_pyd = _find_core_pyd() backup_path = None try: # Windows: proactive backup when the .pyd is accessible (not locked). # If locked, pip will fail before touching the file — just warn. # If not locked, pip will attempt to replace it — back it up as a safety net. if sys.platform.startswith("win") and core_pyd: if _is_pyd_locked(core_pyd): click.secho( "Warning: headroom._core is locked (headroom proxy is running).", fg="yellow", ) click.echo("If the upgrade fails, stop the proxy and try again.") else: click.echo("Creating a backup of headroom._core before upgrading...") try: backup_path = _make_backup(core_pyd) click.echo(f"Backup created: {backup_path}") except Exception as e: raise click.ClickException(f"Could not back up {core_pyd}: {e}") from None # Executing the update result = subprocess.run(argv) # noqa: S603 if result.returncode != 0: # Upgrade failed: restore backup if available if core_pyd and backup_path and backup_path.exists(): click.secho("Restoring backup after upgrade failure...", fg="yellow") try: _restore_backup(core_pyd, backup_path) click.echo("Backup restored successfully.") except Exception as e: click.secho( f"ERROR: Could not restore backup: {e}", fg="red", ) return result.returncode # Upgrade successful: test integrity if core_pyd: click.echo("Testing integrity of headroom._core module...") if not _test_core_integrity(): click.secho( "ERROR: headroom._core could not be imported after upgrade.", fg="red", ) if backup_path and backup_path.exists(): click.secho("Restoring backup...", fg="yellow") try: _restore_backup(core_pyd, backup_path) click.echo("Backup restored successfully.") except Exception as e: click.secho( f"CRITICAL ERROR: Could not restore: {e}", fg="red", ) return 1 click.echo("✓ Module tested successfully.") # Cleanup: remove backup if everything went well if backup_path and backup_path.exists(): try: backup_path.unlink() except Exception: pass # Failed to delete backup is not critical return 0 except FileNotFoundError: if core_pyd and backup_path and backup_path.exists(): try: _restore_backup(core_pyd, backup_path) except Exception: pass raise except click.ClickException: if core_pyd and backup_path and backup_path.exists(): try: _restore_backup(core_pyd, backup_path) except Exception: pass raise def detect_install_method(extras: str | None = None) -> InstallMethod: """Classify the install and build the appropriate upgrade plan. Resolution order (first match wins), covering every supported install path on macOS / Linux / Windows: 1. git checkout → refuse (`git pull`) 2. editable install → refuse (reinstall from source) 3. Docker → refuse (pull a new image) 4. pipx → `pipx upgrade` 5. uv tool → `uv tool upgrade` 6. venv / virtualenv / conda → `sys.executable -m pip install -U` 7. user-site (`pip --user`) → `sys.executable -m pip install -U --user` 8. externally-managed system Python (PEP 668) → refuse with guidance 9. writable global Python → `sys.executable -m pip install -U` (last resort) """ if _is_source_checkout(): return InstallMethod( kind="checkout", can_self_update=False, guidance="Running from a source checkout — update with `git pull`.", ) if _is_editable_install(): return InstallMethod( kind="editable", can_self_update=False, guidance=( "Editable install detected — update your source tree, or " "reinstall with `pip install -U --force-reinstall .`." ), ) if _in_docker(): return InstallMethod( kind="docker", can_self_update=False, guidance=( "Running inside a container — pull a newer Headroom image instead of self-updating." ), ) # pipx / uv tool own their venvs; match against executable, prefix, and the # distribution location so detection works regardless of platform layout. haystack = "::".join( filter( None, ( _norm(getattr(sys, "executable", "")), _norm(getattr(sys, "prefix", "")), _package_location(), ), ) ) pipx_home = _norm(os.environ.get("PIPX_HOME")) if "/pipx/venvs/" in haystack or "/pipx/" in haystack or (pipx_home and pipx_home in haystack): return InstallMethod( kind="pipx", can_self_update=True, argv=["pipx", "upgrade", PACKAGE_NAME], ) uv_tool_dir = _norm(os.environ.get("UV_TOOL_DIR")) if "/uv/tools/" in haystack or (uv_tool_dir and uv_tool_dir in haystack): return InstallMethod( kind="uv-tool", can_self_update=True, argv=["uv", "tool", "upgrade", PACKAGE_NAME], ) if _in_virtualenv(): return InstallMethod( kind="pip", can_self_update=True, argv=[sys.executable, "-m", "pip", "install", "-U", _spec(extras)], ) location = _package_location() if _is_user_site_install(location): return InstallMethod( kind="pip-user", can_self_update=True, argv=[sys.executable, "-m", "pip", "install", "-U", "--user", _spec(extras)], ) if _is_externally_managed(): return InstallMethod( kind="system", can_self_update=False, guidance=_managed_env_guidance(), ) # Writable global interpreter (e.g. Windows python.org, some Linux setups): # pip -U works without admin in the common case; if not, the manual command # is surfaced on failure. return InstallMethod( kind="pip", can_self_update=True, argv=[sys.executable, "-m", "pip", "install", "-U", _spec(extras)], ) @main.command("update") @click.option("--check", "check_only", is_flag=True, help="Report only; do not upgrade.") @click.option("-y", "--yes", "assume_yes", is_flag=True, help="Skip the confirmation prompt.") @click.option("--pre", "allow_pre", is_flag=True, help="Include pre-releases.") @click.option( "--extras", default=None, help="Re-request extras for the pip path, e.g. 'all' or 'proxy'.", ) def update(check_only: bool, assume_yes: bool, allow_pre: bool, extras: str | None) -> None: """Update Headroom to the latest release. Detects pipx / uv tool / pip installs and runs the right upgrade. Refuses (with guidance) for git checkouts, editable installs, Docker, and system Python. """ current = installed_version() click.echo("Checking PyPI for the latest Headroom release...") latest = fetch_latest_version(allow_pre=allow_pre) if latest is None: raise click.ClickException("Could not reach PyPI to check for updates. Try again later.") # Refresh the banner cache as a side effect of an explicit check. write_cache(latest) if current: from packaging.version import InvalidVersion, Version try: if Version(latest) <= Version(current): click.echo(f"Headroom is up to date ({current}).") return except InvalidVersion: pass click.echo(f"Update available: {current} → {latest}") else: click.echo(f"Latest Headroom release: {latest}") method = detect_install_method(extras) if not method.can_self_update: click.echo(method.guidance or "Automatic update is not available for this install.") return assert method.argv is not None cmd_str = _format_cmd(method.argv) click.echo(f"Upgrade command: {cmd_str}") if check_only: return if not assume_yes and not click.confirm("Proceed with the upgrade?", default=True): click.echo("Aborted.") return click.echo(f"Running: {cmd_str}") try: returncode = safe_update(method.argv) except FileNotFoundError: raise click.ClickException( f"`{method.argv[0]}` was not found on PATH. Install it or upgrade manually: {cmd_str}" ) from None if returncode != 0: raise click.ClickException(f"Upgrade failed (exit {returncode}). Run manually: {cmd_str}") # ASCII-only output — emoji can raise UnicodeEncodeError on some Windows consoles. click.echo(f"Headroom upgraded to {latest}.") click.echo("Restart any running `headroom proxy` to pick up the new version.") __all__ = ["InstallMethod", "detect_install_method", "update"]