|
| 1 | +""" |
| 2 | +Auto-update functionality for KiLM. |
| 3 | +Handles installation method detection, PyPI integration, and update execution. |
| 4 | +""" |
| 5 | + |
| 6 | +import json |
| 7 | +import os |
| 8 | +import subprocess |
| 9 | +import sys |
| 10 | +import time |
| 11 | +from pathlib import Path |
| 12 | +from typing import Dict, Optional, Tuple |
| 13 | + |
| 14 | +import requests |
| 15 | +from packaging.version import InvalidVersion, Version |
| 16 | + |
| 17 | + |
| 18 | +def detect_installation_method() -> str: |
| 19 | + """ |
| 20 | + Detect how KiLM was installed to determine appropriate update strategy. |
| 21 | + Returns: 'pipx' | 'pip' | 'pip-venv' | 'uv' | 'conda' |
| 22 | + """ |
| 23 | + executable_path = Path(sys.executable) |
| 24 | + |
| 25 | + # Check for pipx installation |
| 26 | + if any( |
| 27 | + part in str(executable_path) for part in [".local/share/pipx", "pipx/venvs"] |
| 28 | + ): |
| 29 | + return "pipx" |
| 30 | + |
| 31 | + if os.environ.get("PIPX_HOME") and "pipx" in str(executable_path): |
| 32 | + return "pipx" |
| 33 | + |
| 34 | + # Check for conda installation |
| 35 | + if os.environ.get("CONDA_DEFAULT_ENV") or "conda" in str(executable_path): |
| 36 | + return "conda" |
| 37 | + |
| 38 | + # Check for uv installation (only via official environment variables) |
| 39 | + if os.environ.get("UV_TOOL_DIR") or os.environ.get("UV_TOOL_BIN_DIR"): |
| 40 | + return "uv" |
| 41 | + |
| 42 | + # Check for virtual environment (pip in venv) |
| 43 | + if os.environ.get("VIRTUAL_ENV") or sys.prefix != getattr( |
| 44 | + sys, "base_prefix", sys.prefix |
| 45 | + ): |
| 46 | + return "pip-venv" |
| 47 | + |
| 48 | + # Check for homebrew installation (strict path check) |
| 49 | + if str(executable_path).startswith(("/opt/homebrew/", "/usr/local/Cellar/")): |
| 50 | + return "homebrew" |
| 51 | + |
| 52 | + # Default to system pip |
| 53 | + return "pip" |
| 54 | + |
| 55 | + |
| 56 | +class PyPIVersionChecker: |
| 57 | + """Responsible PyPI API client with caching and proper headers.""" |
| 58 | + |
| 59 | + def __init__(self, package_name: str, version: str = "unknown"): |
| 60 | + self.package_name = package_name |
| 61 | + self.base_url = f"https://pypi.org/pypi/{package_name}/json" |
| 62 | + self.cache_file = Path.home() / ".cache" / "kilm" / "version_check.json" |
| 63 | + self.user_agent = f"KiLM/{version} (+https://github.com/barisgit/KiLM)" |
| 64 | + |
| 65 | + def check_latest_version(self) -> Optional[str]: |
| 66 | + """ |
| 67 | + Check latest version from PyPI with caching and rate limiting. |
| 68 | + Returns None if check fails or is rate limited. |
| 69 | + """ |
| 70 | + try: |
| 71 | + headers = {"User-Agent": self.user_agent} |
| 72 | + |
| 73 | + # Use cached ETag if available |
| 74 | + cached_data = self._load_cache() |
| 75 | + if cached_data is not None and "etag" in cached_data: |
| 76 | + headers["If-None-Match"] = cached_data["etag"] |
| 77 | + |
| 78 | + response = requests.get(self.base_url, headers=headers, timeout=10) |
| 79 | + |
| 80 | + if response.status_code == 304: # Not Modified |
| 81 | + return cached_data.get("version") if cached_data else None |
| 82 | + |
| 83 | + if response.status_code == 200: |
| 84 | + data = response.json() |
| 85 | + latest_version = data["info"]["version"] |
| 86 | + |
| 87 | + # Cache response with ETag |
| 88 | + self._save_cache( |
| 89 | + { |
| 90 | + "version": latest_version, |
| 91 | + "etag": response.headers.get("ETag"), |
| 92 | + "timestamp": time.time(), |
| 93 | + } |
| 94 | + ) |
| 95 | + |
| 96 | + return latest_version |
| 97 | + |
| 98 | + except (requests.RequestException, KeyError, json.JSONDecodeError): |
| 99 | + # Fail silently - don't block CLI functionality |
| 100 | + pass |
| 101 | + |
| 102 | + return None |
| 103 | + |
| 104 | + def _load_cache(self) -> Optional[Dict]: |
| 105 | + """Load cached version data.""" |
| 106 | + if self.cache_file.exists(): |
| 107 | + try: |
| 108 | + with Path(self.cache_file).open() as f: |
| 109 | + data = json.load(f) |
| 110 | + # Cache valid for 24 hours |
| 111 | + if time.time() - data.get("timestamp", 0) < 86400: |
| 112 | + return data |
| 113 | + except (json.JSONDecodeError, KeyError): |
| 114 | + pass |
| 115 | + return None |
| 116 | + |
| 117 | + def _save_cache(self, data: Dict): |
| 118 | + """Save version data to cache.""" |
| 119 | + self.cache_file.parent.mkdir(parents=True, exist_ok=True) |
| 120 | + with Path(self.cache_file).open("w") as f: |
| 121 | + json.dump(data, f) |
| 122 | + |
| 123 | + |
| 124 | +def update_via_pipx() -> bool: |
| 125 | + """Update KiLM via pipx. Most reliable method for CLI tools.""" |
| 126 | + try: |
| 127 | + result = subprocess.run( |
| 128 | + ["pipx", "upgrade", "kilm"], capture_output=True, text=True, timeout=300 |
| 129 | + ) |
| 130 | + return result.returncode == 0 |
| 131 | + except (subprocess.TimeoutExpired, FileNotFoundError): |
| 132 | + return False |
| 133 | + |
| 134 | + |
| 135 | +def update_via_pip() -> bool: |
| 136 | + """Update KiLM via pip.""" |
| 137 | + try: |
| 138 | + # Use same Python interpreter that's running KiLM |
| 139 | + result = subprocess.run( |
| 140 | + [sys.executable, "-m", "pip", "install", "--upgrade", "kilm"], |
| 141 | + capture_output=True, |
| 142 | + text=True, |
| 143 | + timeout=300, |
| 144 | + ) |
| 145 | + return result.returncode == 0 |
| 146 | + except (subprocess.TimeoutExpired, FileNotFoundError): |
| 147 | + return False |
| 148 | + |
| 149 | + |
| 150 | +def update_via_uv() -> bool: |
| 151 | + """Update KiLM via uv.""" |
| 152 | + try: |
| 153 | + result = subprocess.run( |
| 154 | + ["uv", "tool", "upgrade", "kilm"], |
| 155 | + capture_output=True, |
| 156 | + text=True, |
| 157 | + timeout=300, |
| 158 | + ) |
| 159 | + return result.returncode == 0 |
| 160 | + except (subprocess.TimeoutExpired, FileNotFoundError): |
| 161 | + return False |
| 162 | + |
| 163 | + |
| 164 | +class UpdateManager: |
| 165 | + """Manages update checking and execution for KiLM.""" |
| 166 | + |
| 167 | + def __init__(self, current_version: str): |
| 168 | + self.version_checker = PyPIVersionChecker("kilm") |
| 169 | + self.current_version = current_version |
| 170 | + self.installation_method = detect_installation_method() |
| 171 | + |
| 172 | + def check_latest_version(self) -> Optional[str]: |
| 173 | + """Check for latest version available on PyPI.""" |
| 174 | + return self.version_checker.check_latest_version() |
| 175 | + |
| 176 | + def is_newer_version_available(self, latest_version: str) -> bool: |
| 177 | + """Compare versions to determine if update is available.""" |
| 178 | + try: |
| 179 | + current_ver = Version(self.current_version) |
| 180 | + latest_ver = Version(latest_version) |
| 181 | + return latest_ver > current_ver |
| 182 | + except (InvalidVersion, AttributeError): |
| 183 | + return False |
| 184 | + |
| 185 | + def get_update_instruction(self) -> str: |
| 186 | + """Get update instruction for the detected installation method.""" |
| 187 | + instructions = { |
| 188 | + "pipx": "pipx upgrade kilm", |
| 189 | + "pip": "pip install --upgrade kilm", |
| 190 | + "pip-venv": "pip install --upgrade kilm", |
| 191 | + "uv": "uv tool upgrade kilm", |
| 192 | + "conda": "Conda package not yet available (planned for future)", |
| 193 | + "homebrew": "Homebrew package not yet available (planned for future)", |
| 194 | + } |
| 195 | + return instructions.get(self.installation_method, "Check your package manager") |
| 196 | + |
| 197 | + def can_auto_update(self) -> bool: |
| 198 | + """Check if automatic update is possible for this installation method.""" |
| 199 | + return self.installation_method in ["pipx", "pip", "pip-venv", "uv"] |
| 200 | + |
| 201 | + def perform_update(self) -> Tuple[bool, str]: |
| 202 | + """ |
| 203 | + Execute update using detected installation method. |
| 204 | + Returns: (success: bool, message: str) |
| 205 | + """ |
| 206 | + if not self.can_auto_update(): |
| 207 | + instruction = self.get_update_instruction() |
| 208 | + return False, f"Manual update required. Run: {instruction}" |
| 209 | + |
| 210 | + update_functions = { |
| 211 | + "pipx": update_via_pipx, |
| 212 | + "pip": update_via_pip, |
| 213 | + "pip-venv": update_via_pip, |
| 214 | + "uv": update_via_uv, |
| 215 | + } |
| 216 | + |
| 217 | + update_func = update_functions.get(self.installation_method) |
| 218 | + if update_func: |
| 219 | + try: |
| 220 | + success = update_func() |
| 221 | + if success: |
| 222 | + return True, "KiLM updated successfully!" |
| 223 | + else: |
| 224 | + instruction = self.get_update_instruction() |
| 225 | + return False, f"Auto-update failed. Try manually: {instruction}" |
| 226 | + except Exception as e: |
| 227 | + return False, f"Update error: {str(e)}" |
| 228 | + else: |
| 229 | + instruction = self.get_update_instruction() |
| 230 | + return False, f"Unsupported installation method. Run: {instruction}" |
0 commit comments