python_check_updates/
global.rs

1use crate::resolver::UpdateSeverity;
2use crate::version::Version;
3use anyhow::Result;
4use std::collections::{HashMap, HashSet};
5use std::fs;
6use std::path::Path;
7use std::process::Command;
8use std::str::FromStr;
9
10/// Source of a globally installed package
11#[derive(Debug, Clone, PartialEq, Eq, Hash)]
12pub enum GlobalSource {
13    Uv,
14    Pipx,
15    PipUser,
16}
17
18impl std::fmt::Display for GlobalSource {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        match self {
21            GlobalSource::Uv => write!(f, "uv"),
22            GlobalSource::Pipx => write!(f, "pipx"),
23            GlobalSource::PipUser => write!(f, "pip"),
24        }
25    }
26}
27
28/// A globally installed package
29#[derive(Debug, Clone)]
30pub struct GlobalPackage {
31    pub name: String,
32    pub installed_version: Version,
33    pub source: GlobalSource,
34    /// Python version (only set for pip --user packages)
35    pub python_version: Option<String>,
36}
37
38/// Result of checking a global package
39#[derive(Debug, Clone)]
40pub struct GlobalCheck {
41    pub package: GlobalPackage,
42    pub latest: Version,
43    pub has_update: bool,
44}
45
46impl GlobalCheck {
47    /// Get update severity for coloring
48    pub fn update_severity(&self) -> Option<UpdateSeverity> {
49        if !self.has_update {
50            return None;
51        }
52        let current = &self.package.installed_version;
53        let target = &self.latest;
54
55        if target.major > current.major {
56            Some(UpdateSeverity::Major)
57        } else if target.minor > current.minor {
58            Some(UpdateSeverity::Minor)
59        } else if target.patch > current.patch {
60            Some(UpdateSeverity::Patch)
61        } else {
62            None
63        }
64    }
65}
66
67/// Discovers globally installed packages from various sources
68pub struct GlobalPackageDiscovery {
69    _include_prerelease: bool,
70}
71
72impl GlobalPackageDiscovery {
73    pub fn new(include_prerelease: bool) -> Self {
74        Self {
75            _include_prerelease: include_prerelease,
76        }
77    }
78
79    /// Discover all globally installed packages
80    pub fn discover(&self) -> Vec<GlobalPackage> {
81        let mut packages = Vec::new();
82
83        // Try each source, silently skip if not available
84        packages.extend(self.discover_uv_tools().unwrap_or_default());
85        packages.extend(self.discover_pipx_packages().unwrap_or_default());
86        packages.extend(self.discover_pip_user_packages().unwrap_or_default());
87
88        packages
89    }
90
91    /// Discover uv tools using `uv tool list`
92    fn discover_uv_tools(&self) -> Result<Vec<GlobalPackage>> {
93        let output = Command::new("uv").args(["tool", "list"]).output();
94
95        match output {
96            Ok(output) if output.status.success() => {
97                self.parse_uv_tool_list(&String::from_utf8_lossy(&output.stdout))
98            }
99            _ => Ok(Vec::new()), // uv not installed or failed, skip silently
100        }
101    }
102
103    /// Parse output of `uv tool list`
104    /// Format: "package_name vX.Y.Z" or "package_name X.Y.Z"
105    /// May also have lines starting with "-" for entry points (skip these)
106    fn parse_uv_tool_list(&self, output: &str) -> Result<Vec<GlobalPackage>> {
107        let mut packages = Vec::new();
108
109        for line in output.lines() {
110            let line = line.trim();
111            if line.is_empty() || line.starts_with('-') {
112                continue;
113            }
114
115            // Parse "name vX.Y.Z" or "name X.Y.Z" format
116            let parts: Vec<&str> = line.split_whitespace().collect();
117            if parts.len() >= 2 {
118                let name = parts[0].to_string();
119                let version_str = parts[1].trim_start_matches('v');
120
121                if let Ok(version) = Version::from_str(version_str) {
122                    packages.push(GlobalPackage {
123                        name,
124                        installed_version: version,
125                        source: GlobalSource::Uv,
126                        python_version: None,
127                    });
128                }
129            }
130        }
131
132        Ok(packages)
133    }
134
135    /// Discover pipx packages
136    fn discover_pipx_packages(&self) -> Result<Vec<GlobalPackage>> {
137        // Try `pipx list --json` for structured output
138        let output = Command::new("pipx").args(["list", "--json"]).output();
139
140        match output {
141            Ok(output) if output.status.success() => {
142                self.parse_pipx_json(&String::from_utf8_lossy(&output.stdout))
143            }
144            _ => {
145                // Fall back to scanning ~/.local/pipx/venvs/
146                self.discover_pipx_from_directory()
147            }
148        }
149    }
150
151    /// Parse pipx list --json output
152    fn parse_pipx_json(&self, json_str: &str) -> Result<Vec<GlobalPackage>> {
153        let data: serde_json::Value = serde_json::from_str(json_str)?;
154        let mut packages = Vec::new();
155
156        if let Some(venvs) = data.get("venvs").and_then(|v| v.as_object()) {
157            for (name, venv_data) in venvs {
158                if let Some(version_str) = venv_data
159                    .pointer("/metadata/main_package/package_version")
160                    .and_then(|v| v.as_str())
161                {
162                    if let Ok(version) = Version::from_str(version_str) {
163                        packages.push(GlobalPackage {
164                            name: name.clone(),
165                            installed_version: version,
166                            source: GlobalSource::Pipx,
167                            python_version: None,
168                        });
169                    }
170                }
171            }
172        }
173
174        Ok(packages)
175    }
176
177    /// Fall back: scan ~/.local/pipx/venvs/ directory
178    fn discover_pipx_from_directory(&self) -> Result<Vec<GlobalPackage>> {
179        let pipx_dir = dirs::home_dir()
180            .map(|h| h.join(".local/pipx/venvs"))
181            .filter(|p| p.exists());
182
183        let Some(pipx_dir) = pipx_dir else {
184            return Ok(Vec::new());
185        };
186
187        let mut packages = Vec::new();
188
189        for entry in fs::read_dir(&pipx_dir)? {
190            let entry = entry?;
191            if entry.path().is_dir() {
192                let name = entry.file_name().to_string_lossy().to_string();
193
194                // Try to get version from the package's metadata
195                if let Some(version) = self.get_pipx_package_version(&entry.path(), &name) {
196                    packages.push(GlobalPackage {
197                        name,
198                        installed_version: version,
199                        source: GlobalSource::Pipx,
200                        python_version: None,
201                    });
202                }
203            }
204        }
205
206        Ok(packages)
207    }
208
209    /// Get version of a pipx package by reading its dist-info
210    fn get_pipx_package_version(&self, venv_path: &Path, package_name: &str) -> Option<Version> {
211        // Look in the venv's site-packages for the dist-info
212        let site_packages = venv_path.join("lib");
213
214        if !site_packages.exists() {
215            return None;
216        }
217
218        // Find the python directory (e.g., python3.11)
219        let python_dir = fs::read_dir(&site_packages)
220            .ok()?
221            .filter_map(|e| e.ok())
222            .find(|e| e.file_name().to_string_lossy().starts_with("python"))?;
223
224        let actual_site_packages = python_dir.path().join("site-packages");
225        if !actual_site_packages.exists() {
226            return None;
227        }
228
229        // Look for the dist-info directory
230        let normalized_name = package_name.to_lowercase().replace('-', "_");
231        for entry in fs::read_dir(&actual_site_packages).ok()? {
232            let entry = entry.ok()?;
233            let name = entry.file_name().to_string_lossy().to_string();
234            if name.ends_with(".dist-info") {
235                let dist_name = name
236                    .strip_suffix(".dist-info")?
237                    .to_lowercase()
238                    .replace('-', "_");
239                // Check if this dist-info matches our package
240                if dist_name.starts_with(&normalized_name) {
241                    if let Some((_, version)) = self.parse_dist_info_name(&name) {
242                        return Some(version);
243                    }
244                }
245            }
246        }
247
248        None
249    }
250
251    /// Discover pip --user packages from ~/.local/lib/python3.x/site-packages/
252    /// Now tracks which Python version each package belongs to
253    fn discover_pip_user_packages(&self) -> Result<Vec<GlobalPackage>> {
254        let user_lib = dirs::home_dir().map(|h| h.join(".local/lib"));
255
256        let Some(user_lib) = user_lib else {
257            return Ok(Vec::new());
258        };
259
260        if !user_lib.exists() {
261            return Ok(Vec::new());
262        }
263
264        let mut packages = Vec::new();
265
266        // Find all python3.x directories and collect them sorted
267        let mut python_dirs: Vec<_> = fs::read_dir(&user_lib)?
268            .filter_map(|e| e.ok())
269            .filter(|e| {
270                let name = e.file_name().to_string_lossy().to_string();
271                name.starts_with("python3.") || name.starts_with("python2.")
272            })
273            .collect();
274
275        // Sort by version descending so newer Python versions come first
276        python_dirs.sort_by(|a, b| {
277            let a_name = a.file_name().to_string_lossy().to_string();
278            let b_name = b.file_name().to_string_lossy().to_string();
279            b_name.cmp(&a_name)
280        });
281
282        // Track seen packages to avoid duplicates across Python versions
283        let mut seen_packages: HashSet<String> = HashSet::new();
284
285        for entry in python_dirs {
286            let dir_name = entry.file_name().to_string_lossy().to_string();
287            // Extract version like "3.11" from "python3.11"
288            let python_version = dir_name.strip_prefix("python").unwrap_or(&dir_name);
289
290            let site_packages = entry.path().join("site-packages");
291            if site_packages.exists() {
292                packages.extend(self.parse_site_packages(
293                    &site_packages,
294                    python_version,
295                    &mut seen_packages,
296                )?);
297            }
298        }
299
300        Ok(packages)
301    }
302
303    /// Parse a site-packages directory for installed packages
304    fn parse_site_packages(
305        &self,
306        site_packages: &Path,
307        python_version: &str,
308        seen: &mut HashSet<String>,
309    ) -> Result<Vec<GlobalPackage>> {
310        let mut packages = Vec::new();
311
312        // Look for .dist-info directories
313        for entry in fs::read_dir(site_packages)? {
314            let entry = entry?;
315            let name = entry.file_name().to_string_lossy().to_string();
316
317            if name.ends_with(".dist-info") {
318                // Parse: "package_name-1.2.3.dist-info"
319                if let Some((pkg_name, version)) = self.parse_dist_info_name(&name) {
320                    // Normalize package name for deduplication
321                    let normalized = pkg_name.to_lowercase().replace('-', "_");
322
323                    // Skip if we've already seen this package (from another Python version)
324                    if seen.contains(&normalized) {
325                        continue;
326                    }
327                    seen.insert(normalized);
328
329                    packages.push(GlobalPackage {
330                        name: pkg_name,
331                        installed_version: version,
332                        source: GlobalSource::PipUser,
333                        python_version: Some(python_version.to_string()),
334                    });
335                }
336            }
337        }
338
339        Ok(packages)
340    }
341
342    /// Parse a dist-info directory name to extract package name and version
343    /// Format: "package_name-1.2.3.dist-info"
344    fn parse_dist_info_name(&self, name: &str) -> Option<(String, Version)> {
345        let without_suffix = name.strip_suffix(".dist-info")?;
346
347        // Find the last hyphen that separates name from version
348        // Version always starts with a digit
349        let mut split_idx = None;
350        for (i, c) in without_suffix.char_indices().rev() {
351            if c == '-' {
352                // Check if what follows is a version (starts with digit)
353                if without_suffix[i + 1..]
354                    .chars()
355                    .next()
356                    .is_some_and(|c| c.is_ascii_digit())
357                {
358                    split_idx = Some(i);
359                    break;
360                }
361            }
362        }
363
364        let idx = split_idx?;
365        let pkg_name = &without_suffix[..idx];
366        let version_str = &without_suffix[idx + 1..];
367
368        let version = Version::from_str(version_str).ok()?;
369        Some((pkg_name.to_string(), version))
370    }
371}
372
373/// Group checks by source for upgrade command generation
374pub fn group_by_source(checks: &[GlobalCheck]) -> HashMap<GlobalSource, Vec<&GlobalCheck>> {
375    checks
376        .iter()
377        .filter(|c| c.has_update)
378        .fold(HashMap::new(), |mut acc, check| {
379            acc.entry(check.package.source.clone())
380                .or_insert_with(Vec::new)
381                .push(check);
382            acc
383        })
384}
385
386/// Check if a Python version is available on the system
387pub fn is_python_available(version: &str) -> bool {
388    // Try python3.X --version
389    let cmd = format!("python{}", version);
390    Command::new(&cmd)
391        .arg("--version")
392        .output()
393        .map(|o| o.status.success())
394        .unwrap_or(false)
395}
396
397/// Get the user's home directory path for display
398fn get_pip_user_path(python_version: &str) -> String {
399    dirs::home_dir()
400        .map(|h| h.join(format!(".local/lib/python{}", python_version)))
401        .map(|p| p.display().to_string())
402        .unwrap_or_else(|| format!("~/.local/lib/python{}", python_version))
403}
404
405/// An upgrade command or a comment (for unavailable Python versions)
406#[derive(Debug, Clone)]
407pub enum UpgradeCommand {
408    /// A shell command to run
409    Command(String),
410    /// A comment (Python version not available)
411    Comment(String),
412}
413
414/// Generate upgrade commands for each source
415pub fn generate_upgrade_commands(checks: &[GlobalCheck]) -> Vec<UpgradeCommand> {
416    let updates_by_source = group_by_source(checks);
417    let mut commands = Vec::new();
418
419    if updates_by_source.contains_key(&GlobalSource::Uv) {
420        commands.push(UpgradeCommand::Command("uv tool upgrade --all".to_string()));
421    }
422
423    if updates_by_source.contains_key(&GlobalSource::Pipx) {
424        commands.push(UpgradeCommand::Command("pipx upgrade-all".to_string()));
425    }
426
427    if let Some(pip_updates) = updates_by_source.get(&GlobalSource::PipUser) {
428        // Group pip packages by Python version
429        let mut by_python: std::collections::BTreeMap<String, Vec<&str>> =
430            std::collections::BTreeMap::new();
431
432        for check in pip_updates {
433            let py_version = check
434                .package
435                .python_version
436                .clone()
437                .unwrap_or_else(|| "unknown".to_string());
438            by_python
439                .entry(py_version)
440                .or_insert_with(Vec::new)
441                .push(check.package.name.as_str());
442        }
443
444        // Generate a command for each Python version
445        for (py_version, package_names) in by_python {
446            if is_python_available(&py_version) {
447                commands.push(UpgradeCommand::Command(format!(
448                    "python{} -m pip install --user --upgrade {}",
449                    py_version,
450                    package_names.join(" ")
451                )));
452            } else {
453                let path = get_pip_user_path(&py_version);
454                commands.push(UpgradeCommand::Comment(format!(
455                    "Python {} is no longer installed. Consider removing {} if nothing uses it.",
456                    py_version, path
457                )));
458            }
459        }
460    }
461
462    commands
463}
464
465#[cfg(test)]
466mod tests {
467    use super::*;
468
469    #[test]
470    fn test_parse_uv_tool_list() {
471        let discovery = GlobalPackageDiscovery::new(false);
472        let output = r#"ruff v0.14.10
473    - ruff
474ty v0.0.5
475    - ty
476"#;
477        let packages = discovery.parse_uv_tool_list(output).unwrap();
478        assert_eq!(packages.len(), 2);
479        assert_eq!(packages[0].name, "ruff");
480        assert_eq!(packages[0].installed_version.to_string(), "0.14.10");
481        assert_eq!(packages[0].source, GlobalSource::Uv);
482        assert_eq!(packages[1].name, "ty");
483        assert_eq!(packages[1].installed_version.to_string(), "0.0.5");
484    }
485
486    #[test]
487    fn test_parse_uv_tool_list_without_v_prefix() {
488        let discovery = GlobalPackageDiscovery::new(false);
489        let output = "black 24.10.0\n";
490        let packages = discovery.parse_uv_tool_list(output).unwrap();
491        assert_eq!(packages.len(), 1);
492        assert_eq!(packages[0].name, "black");
493        assert_eq!(packages[0].installed_version.to_string(), "24.10.0");
494    }
495
496    #[test]
497    fn test_parse_pipx_json() {
498        let discovery = GlobalPackageDiscovery::new(false);
499        let json = r#"{
500            "venvs": {
501                "black": {
502                    "metadata": {
503                        "main_package": {
504                            "package_version": "24.10.0"
505                        }
506                    }
507                },
508                "ruff": {
509                    "metadata": {
510                        "main_package": {
511                            "package_version": "0.14.9"
512                        }
513                    }
514                }
515            }
516        }"#;
517        let packages = discovery.parse_pipx_json(json).unwrap();
518        assert_eq!(packages.len(), 2);
519        // Note: HashMap iteration order is not guaranteed, so we check by finding
520        let black = packages.iter().find(|p| p.name == "black").unwrap();
521        assert_eq!(black.installed_version.to_string(), "24.10.0");
522        assert_eq!(black.source, GlobalSource::Pipx);
523    }
524
525    #[test]
526    fn test_parse_dist_info_name() {
527        let discovery = GlobalPackageDiscovery::new(false);
528
529        // Simple case
530        let result = discovery.parse_dist_info_name("requests-2.28.0.dist-info");
531        assert!(result.is_some());
532        let (name, version) = result.unwrap();
533        assert_eq!(name, "requests");
534        assert_eq!(version.to_string(), "2.28.0");
535
536        // Package with hyphen in name
537        let result = discovery.parse_dist_info_name("typing-extensions-4.12.2.dist-info");
538        assert!(result.is_some());
539        let (name, version) = result.unwrap();
540        assert_eq!(name, "typing-extensions");
541        assert_eq!(version.to_string(), "4.12.2");
542
543        // Package with underscore
544        let result = discovery.parse_dist_info_name("my_package-1.0.0.dist-info");
545        assert!(result.is_some());
546        let (name, version) = result.unwrap();
547        assert_eq!(name, "my_package");
548        assert_eq!(version.to_string(), "1.0.0");
549    }
550
551    #[test]
552    fn test_global_source_display() {
553        assert_eq!(GlobalSource::Uv.to_string(), "uv");
554        assert_eq!(GlobalSource::Pipx.to_string(), "pipx");
555        assert_eq!(GlobalSource::PipUser.to_string(), "pip");
556    }
557
558    #[test]
559    fn test_update_severity() {
560        let pkg = GlobalPackage {
561            name: "test".to_string(),
562            installed_version: Version::from_str("1.0.0").unwrap(),
563            source: GlobalSource::Uv,
564            python_version: None,
565        };
566
567        // Major update
568        let check = GlobalCheck {
569            package: pkg.clone(),
570            latest: Version::from_str("2.0.0").unwrap(),
571            has_update: true,
572        };
573        assert_eq!(check.update_severity(), Some(UpdateSeverity::Major));
574
575        // Minor update
576        let check = GlobalCheck {
577            package: pkg.clone(),
578            latest: Version::from_str("1.1.0").unwrap(),
579            has_update: true,
580        };
581        assert_eq!(check.update_severity(), Some(UpdateSeverity::Minor));
582
583        // Patch update
584        let check = GlobalCheck {
585            package: pkg.clone(),
586            latest: Version::from_str("1.0.1").unwrap(),
587            has_update: true,
588        };
589        assert_eq!(check.update_severity(), Some(UpdateSeverity::Patch));
590
591        // No update
592        let check = GlobalCheck {
593            package: pkg,
594            latest: Version::from_str("1.0.0").unwrap(),
595            has_update: false,
596        };
597        assert_eq!(check.update_severity(), None);
598    }
599
600    #[test]
601    fn test_generate_upgrade_commands() {
602        let checks = vec![
603            GlobalCheck {
604                package: GlobalPackage {
605                    name: "ruff".to_string(),
606                    installed_version: Version::from_str("0.14.9").unwrap(),
607                    source: GlobalSource::Uv,
608                    python_version: None,
609                },
610                latest: Version::from_str("0.14.10").unwrap(),
611                has_update: true,
612            },
613            GlobalCheck {
614                package: GlobalPackage {
615                    name: "black".to_string(),
616                    installed_version: Version::from_str("24.1.0").unwrap(),
617                    source: GlobalSource::Pipx,
618                    python_version: None,
619                },
620                latest: Version::from_str("24.10.0").unwrap(),
621                has_update: true,
622            },
623            GlobalCheck {
624                package: GlobalPackage {
625                    name: "requests".to_string(),
626                    installed_version: Version::from_str("2.28.0").unwrap(),
627                    source: GlobalSource::PipUser,
628                    python_version: Some("3.11".to_string()),
629                },
630                latest: Version::from_str("2.32.3").unwrap(),
631                has_update: true,
632            },
633            GlobalCheck {
634                package: GlobalPackage {
635                    name: "flask".to_string(),
636                    installed_version: Version::from_str("2.3.3").unwrap(),
637                    source: GlobalSource::PipUser,
638                    python_version: Some("3.11".to_string()),
639                },
640                latest: Version::from_str("3.0.0").unwrap(),
641                has_update: true,
642            },
643        ];
644
645        let commands = generate_upgrade_commands(&checks);
646        // Should have uv, pipx, and either a pip command or comment for 3.11
647        assert!(commands.len() >= 3);
648
649        // Check for uv command
650        let has_uv = commands.iter().any(|c| matches!(c, UpgradeCommand::Command(s) if s == "uv tool upgrade --all"));
651        assert!(has_uv, "Should have uv upgrade command");
652
653        // Check for pipx command
654        let has_pipx = commands.iter().any(|c| matches!(c, UpgradeCommand::Command(s) if s == "pipx upgrade-all"));
655        assert!(has_pipx, "Should have pipx upgrade command");
656
657        // Check for pip command or comment for Python 3.11
658        let has_pip_311 = commands.iter().any(|c| {
659            match c {
660                UpgradeCommand::Command(s) => s.contains("python3.11") && s.contains("requests") && s.contains("flask"),
661                UpgradeCommand::Comment(s) => s.contains("3.11"),
662            }
663        });
664        assert!(has_pip_311, "Should have pip command or comment for Python 3.11");
665    }
666}