rusty_stack/installers/common/
utils.rs1use std::fmt;
13use std::io::IsTerminal;
14use std::path::Path;
15use std::process::Command;
16
17#[derive(Debug, Clone, Copy)]
29pub struct Colors {
30 pub red: &'static str,
31 pub green: &'static str,
32 pub yellow: &'static str,
33 pub blue: &'static str,
34 pub magenta: &'static str,
35 pub cyan: &'static str,
36 pub bold: &'static str,
37 pub reset: &'static str,
38}
39
40const COLOR_ENABLED: Colors = Colors {
42 red: "\x1b[0;31m",
43 green: "\x1b[0;32m",
44 yellow: "\x1b[0;33m",
45 blue: "\x1b[0;34m",
46 magenta: "\x1b[0;35m",
47 cyan: "\x1b[0;36m",
48 bold: "\x1b[1m",
49 reset: "\x1b[0m",
50};
51
52const COLOR_DISABLED: Colors = Colors {
54 red: "",
55 green: "",
56 yellow: "",
57 blue: "",
58 magenta: "",
59 cyan: "",
60 bold: "",
61 reset: "",
62};
63
64pub fn color_enabled() -> bool {
70 if std::env::var("NO_COLOR").is_ok() {
71 return false;
72 }
73 std::io::stderr().is_terminal()
74}
75
76pub fn get_colors() -> &'static Colors {
78 if color_enabled() {
79 &COLOR_ENABLED
80 } else {
81 &COLOR_DISABLED
82 }
83}
84
85pub fn command_exists(cmd: &str) -> bool {
108 Command::new("which")
109 .arg(cmd)
110 .output()
111 .map(|o| o.status.success())
112 .unwrap_or(false)
113}
114
115pub fn resolve_rocminfo_path() -> String {
128 let standard = Path::new("/opt/rocm/bin/rocminfo");
130 if standard.exists() {
131 return standard.to_string_lossy().to_string();
132 }
133
134 if let Ok(entries) = std::fs::read_dir("/opt") {
136 let mut rocm_dirs: Vec<_> = entries
137 .filter_map(|e| e.ok())
138 .filter(|e| e.file_name().to_string_lossy().starts_with("rocm-"))
139 .collect();
140 rocm_dirs.sort_by_key(|b| std::cmp::Reverse(b.file_name()));
141 for dir in rocm_dirs {
142 let path = dir.path().join("bin/rocminfo");
143 if path.exists() {
144 return path.to_string_lossy().to_string();
145 }
146 }
147 }
148
149 let arch_path = Path::new("/usr/lib/rocm/bin/rocminfo");
151 if arch_path.exists() {
152 return arch_path.to_string_lossy().to_string();
153 }
154
155 "rocminfo".to_string()
157}
158
159pub fn resolve_rocprofiler_tool_path() -> Option<String> {
168 let candidates = [
169 "/opt/rocm/lib/rocprofiler-sdk/librocprofiler-sdk-tool.so",
170 "/opt/rocm/lib/librocprofiler-sdk-tool.so",
171 "/usr/lib/rocm/lib/rocprofiler-sdk/librocprofiler-sdk-tool.so",
172 "/usr/lib/rocm/lib/librocprofiler-sdk-tool.so",
173 ];
174 for candidate in &candidates {
175 if Path::new(candidate).exists() {
176 return Some(candidate.to_string());
177 }
178 }
179
180 if let Ok(entries) = std::fs::read_dir("/opt") {
182 let mut rocm_dirs: Vec<_> = entries
183 .filter_map(|e| e.ok())
184 .filter(|e| e.file_name().to_string_lossy().starts_with("rocm-"))
185 .collect();
186 rocm_dirs.sort_by_key(|b| std::cmp::Reverse(b.file_name()));
187 for dir in rocm_dirs {
188 let base = dir.path();
189 for subpath in [
190 "lib/rocprofiler-sdk/librocprofiler-sdk-tool.so",
191 "lib/librocprofiler-sdk-tool.so",
192 ] {
193 let path = base.join(subpath);
194 if path.exists() {
195 return Some(path.to_string_lossy().to_string());
196 }
197 }
198 }
199 }
200
201 None
202}
203
204pub fn print_header(title: &str) {
213 let c = get_colors();
214 let width = 57;
215 let inner = width - 2; let padded = format!("=== {title} ===");
217 let centered = center_str(&padded, inner);
218
219 let border_line = "═".repeat(inner);
220 let empty_line = " ".repeat(inner);
221
222 eprintln!();
223 eprintln!("{}╔{border_line}╗{}", c.cyan, c.reset);
224 eprintln!("{}║{empty_line}║{}", c.cyan, c.reset);
225 eprintln!("{}║{centered}║{}", c.cyan, c.reset);
226 eprintln!("{}║{empty_line}║{}", c.cyan, c.reset);
227 eprintln!("{}╚{border_line}╝{}", c.cyan, c.reset);
228 eprintln!();
229}
230
231pub fn print_section(title: &str) {
233 let c = get_colors();
234 eprintln!();
235 eprintln!(
236 "{}┌─────────────────────────────────────────────────────────┐{}",
237 c.blue, c.reset
238 );
239 eprintln!("{}│ {}{}{}", c.blue, c.bold, title, c.reset);
240 eprintln!(
241 "{}└─────────────────────────────────────────────────────────┘{}",
242 c.blue, c.reset
243 );
244}
245
246pub fn print_step(message: &str) {
248 let c = get_colors();
249 eprintln!("{}➤ {}{}", c.magenta, message, c.reset);
250}
251
252pub fn print_success(message: &str) {
254 let c = get_colors();
255 eprintln!("{}✓ {}{}", c.green, message, c.reset);
256}
257
258pub fn print_warning(message: &str) {
260 let c = get_colors();
261 eprintln!("{}⚠ {}{}", c.yellow, message, c.reset);
262}
263
264pub fn print_error(message: &str) {
266 let c = get_colors();
267 eprintln!("{}✗ {}{}", c.red, message, c.reset);
268}
269
270pub fn execute_command(cmd: &str, description: &str, dry_run: bool) -> Result<(), String> {
275 let c = get_colors();
276 if dry_run {
277 eprintln!("{}[DRY-RUN]{} {}: ", c.yellow, c.reset, description);
278 eprintln!(" {}{}{}", c.bold, cmd, c.reset);
279 return Ok(());
280 }
281
282 print_step(&format!("{description}..."));
283
284 let result = Command::new("bash").arg("-c").arg(cmd).output();
285
286 match result {
287 Ok(output) if output.status.success() => {
288 print_success(&format!("Done: {description}"));
289 Ok(())
290 }
291 Ok(output) => {
292 let stderr = String::from_utf8_lossy(&output.stderr);
293 print_error(&format!("Failed: {description}"));
294 if !stderr.is_empty() {
295 eprintln!(" {stderr}");
296 }
297 Err(format!("Command failed: {cmd}"))
298 }
299 Err(e) => {
300 print_error(&format!("Failed: {description}: {e}"));
301 Err(format!("Command execution error: {e}"))
302 }
303 }
304}
305
306#[derive(Debug, Clone, Copy, PartialEq, Eq)]
312pub enum PythonPkgManager {
313 Uv,
315 Pip,
317}
318
319impl fmt::Display for PythonPkgManager {
320 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
321 match self {
322 PythonPkgManager::Uv => write!(f, "uv"),
323 PythonPkgManager::Pip => write!(f, "pip"),
324 }
325 }
326}
327
328pub fn detect_python_pkg_manager() -> PythonPkgManager {
337 if command_exists("uv") {
338 PythonPkgManager::Uv
339 } else {
340 PythonPkgManager::Pip
341 }
342}
343
344pub fn pip_install_prefix(python_bin: &str) -> Vec<String> {
348 let mgr = detect_python_pkg_manager();
349 match mgr {
350 PythonPkgManager::Uv => {
351 let mut args = vec!["uv".to_string(), "pip".to_string(), "install".to_string()];
352 if std::env::var("UV_PYTHON").is_err() {
354 args.push("--python".to_string());
355 args.push(python_bin.to_string());
356 }
357 args
358 }
359 PythonPkgManager::Pip => {
360 vec![
361 python_bin.to_string(),
362 "-m".to_string(),
363 "pip".to_string(),
364 "install".to_string(),
365 ]
366 }
367 }
368}
369
370#[derive(Debug, Clone, PartialEq, Eq)]
376pub enum PipCacheFixResult {
377 NoCacheDir,
379 OwnershipCorrect,
381 Fixed,
383 FixFailed(String),
385}
386
387impl std::fmt::Display for PipCacheFixResult {
388 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
389 match self {
390 PipCacheFixResult::NoCacheDir => write!(f, "No pip cache directory found"),
391 PipCacheFixResult::OwnershipCorrect => write!(f, "Pip cache ownership already correct"),
392 PipCacheFixResult::Fixed => write!(f, "Fixed pip cache ownership"),
393 PipCacheFixResult::FixFailed(e) => write!(f, "Failed to fix pip cache ownership: {e}"),
394 }
395 }
396}
397
398#[cfg(unix)]
425pub fn fix_pip_cache_ownership(user_home: &str) -> PipCacheFixResult {
426 use std::os::unix::fs::MetadataExt;
427
428 let cache_dir = Path::new(user_home).join(".cache/pip");
429
430 if !cache_dir.exists() {
432 return PipCacheFixResult::NoCacheDir;
433 }
434
435 let home_meta = match std::fs::metadata(user_home) {
439 Ok(m) => m,
440 Err(_) => return PipCacheFixResult::FixFailed("Cannot stat home directory".into()),
441 };
442 let target_uid = home_meta.uid();
443 let target_gid = home_meta.gid();
444
445 let cache_meta = match std::fs::metadata(&cache_dir) {
447 Ok(m) => m,
448 Err(_) => return PipCacheFixResult::FixFailed("Cannot stat cache directory".into()),
449 };
450
451 if cache_meta.uid() == target_uid && cache_meta.gid() == target_gid {
453 return PipCacheFixResult::OwnershipCorrect;
454 }
455
456 let chown_cmd = format!(
460 "chown -R {}:{} \"{}\"",
461 target_uid,
462 target_gid,
463 cache_dir.display()
464 );
465
466 match Command::new("bash").arg("-c").arg(&chown_cmd).output() {
467 Ok(output) if output.status.success() => PipCacheFixResult::Fixed,
468 Ok(output) => {
469 let stderr = String::from_utf8_lossy(&output.stderr);
470 PipCacheFixResult::FixFailed(format!("chown failed: {stderr}"))
471 }
472 Err(e) => PipCacheFixResult::FixFailed(format!("chown execution error: {e}")),
473 }
474}
475
476#[cfg(not(unix))]
480pub fn fix_pip_cache_ownership(_user_home: &str) -> PipCacheFixResult {
481 PipCacheFixResult::OwnershipCorrect
482}
483
484fn center_str(s: &str, width: usize) -> String {
490 if s.len() >= width {
491 return s.to_string();
492 }
493 let left = (width - s.len()) / 2;
494 let right = width - s.len() - left;
495 format!("{}{}{}", " ".repeat(left), s, " ".repeat(right))
496}
497
498#[cfg(test)]
503mod tests {
504 use super::*;
505
506 #[test]
509 fn test_colors_enabled_has_escapes() {
510 let c = &COLOR_ENABLED;
511 assert!(!c.red.is_empty());
512 assert!(!c.green.is_empty());
513 assert!(!c.reset.is_empty());
514 assert!(c.red.starts_with('\x1b'));
515 }
516
517 #[test]
518 fn test_colors_disabled_is_empty() {
519 let c = &COLOR_DISABLED;
520 assert!(c.red.is_empty());
521 assert!(c.green.is_empty());
522 assert!(c.yellow.is_empty());
523 assert!(c.blue.is_empty());
524 assert!(c.magenta.is_empty());
525 assert!(c.cyan.is_empty());
526 assert!(c.bold.is_empty());
527 assert!(c.reset.is_empty());
528 }
529
530 #[test]
531 fn test_get_colors_returns_valid_ref() {
532 let _colors = get_colors();
533 }
535
536 #[test]
539 fn test_command_exists_true() {
540 assert!(command_exists("bash"));
542 }
543
544 #[test]
545 fn test_command_exists_false() {
546 assert!(!command_exists("nonexistent_command_xyz_12345"));
547 }
548
549 #[test]
550 fn test_command_exists_which() {
551 assert!(command_exists("which"));
553 }
554
555 #[test]
558 fn test_print_header_no_panic() {
559 print_header("Test Title");
561 }
562
563 #[test]
564 fn test_print_section_no_panic() {
565 print_section("Test Section");
566 }
567
568 #[test]
569 fn test_print_step_no_panic() {
570 print_step("Test step");
571 }
572
573 #[test]
574 fn test_print_success_no_panic() {
575 print_success("Test success");
576 }
577
578 #[test]
579 fn test_print_warning_no_panic() {
580 print_warning("Test warning");
581 }
582
583 #[test]
584 fn test_print_error_no_panic() {
585 print_error("Test error");
586 }
587
588 #[test]
591 fn test_execute_command_dry_run() {
592 let result = execute_command("echo hello", "test dry run", true);
593 assert!(result.is_ok());
594 }
595
596 #[test]
597 fn test_execute_command_success() {
598 let result = execute_command("true", "test success", false);
599 assert!(result.is_ok());
600 }
601
602 #[test]
603 fn test_execute_command_failure() {
604 let result = execute_command("false", "test failure", false);
605 assert!(result.is_err());
606 }
607
608 #[test]
611 fn test_detect_python_pkg_manager() {
612 let mgr = detect_python_pkg_manager();
614 assert!(mgr == PythonPkgManager::Uv || mgr == PythonPkgManager::Pip);
615 }
616
617 #[test]
618 fn test_python_pkg_manager_display() {
619 assert_eq!(format!("{}", PythonPkgManager::Uv), "uv");
620 assert_eq!(format!("{}", PythonPkgManager::Pip), "pip");
621 }
622
623 #[test]
624 fn test_pip_install_prefix_returns_nonempty() {
625 let prefix = pip_install_prefix("python3");
626 assert!(!prefix.is_empty());
627 assert!(prefix[0] == "uv" || prefix[0] == "python3");
629 }
630
631 #[test]
634 fn test_center_str_short() {
635 let result = center_str("hi", 10);
636 assert_eq!(result.len(), 10);
637 assert_eq!(result.trim(), "hi");
638 }
639
640 #[test]
641 fn test_center_str_exact() {
642 let result = center_str("hello", 5);
643 assert_eq!(result, "hello");
644 }
645
646 #[test]
647 fn test_center_str_longer() {
648 let result = center_str("hello world", 5);
649 assert_eq!(result, "hello world");
650 }
651
652 #[test]
655 fn test_resolve_rocminfo_path_returns_string() {
656 let path = resolve_rocminfo_path();
658 assert!(
659 !path.is_empty(),
660 "resolve_rocminfo_path must return non-empty string"
661 );
662 }
663
664 #[test]
665 fn test_resolve_rocminfo_path_prefers_standard_path() {
666 let path = resolve_rocminfo_path();
668 if Path::new("/opt/rocm/bin/rocminfo").exists() {
669 assert_eq!(path, "/opt/rocm/bin/rocminfo");
670 }
671 assert!(!path.is_empty());
673 }
674
675 #[test]
676 fn test_resolve_rocminfo_path_fallback_is_rocminfo() {
677 let path = resolve_rocminfo_path();
681 assert!(
682 path == "rocminfo" || path.contains("rocminfo"),
683 "Path should contain 'rocminfo', got: {path}"
684 );
685 }
686
687 #[test]
690 fn test_pip_cache_fix_result_display_no_cache_dir() {
691 assert_eq!(
692 format!("{}", PipCacheFixResult::NoCacheDir),
693 "No pip cache directory found"
694 );
695 }
696
697 #[test]
698 fn test_pip_cache_fix_result_display_ownership_correct() {
699 assert_eq!(
700 format!("{}", PipCacheFixResult::OwnershipCorrect),
701 "Pip cache ownership already correct"
702 );
703 }
704
705 #[test]
706 fn test_pip_cache_fix_result_display_fixed() {
707 assert_eq!(
708 format!("{}", PipCacheFixResult::Fixed),
709 "Fixed pip cache ownership"
710 );
711 }
712
713 #[test]
714 fn test_pip_cache_fix_result_display_fix_failed() {
715 let result = PipCacheFixResult::FixFailed("some error".into());
716 assert!(format!("{}", result).contains("some error"));
717 }
718
719 #[test]
720 fn test_pip_cache_fix_result_equality() {
721 assert_eq!(PipCacheFixResult::NoCacheDir, PipCacheFixResult::NoCacheDir);
722 assert_eq!(
723 PipCacheFixResult::OwnershipCorrect,
724 PipCacheFixResult::OwnershipCorrect
725 );
726 assert_eq!(PipCacheFixResult::Fixed, PipCacheFixResult::Fixed);
727 assert_ne!(PipCacheFixResult::NoCacheDir, PipCacheFixResult::Fixed);
728 }
729
730 #[test]
731 fn test_fix_pip_cache_no_cache_dir() {
732 let result = fix_pip_cache_ownership("/tmp/nonexistent_test_dir_xyz");
734 assert_eq!(result, PipCacheFixResult::NoCacheDir);
735 }
736
737 #[test]
738 fn test_fix_pip_cache_ownership_correct() {
739 let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".into());
743 let result = fix_pip_cache_ownership(&home);
744 assert!(
747 result == PipCacheFixResult::NoCacheDir
748 || result == PipCacheFixResult::OwnershipCorrect
749 || result == PipCacheFixResult::Fixed,
750 "Expected NoCacheDir, OwnershipCorrect, or Fixed, got: {result}"
751 );
752 }
753
754 #[test]
755 fn test_fix_pip_cache_with_temp_dir_correct_ownership() {
756 let tmp_dir = std::env::temp_dir().join("rusty_stack_test_pip_cache_correct");
758 let cache_dir = tmp_dir.join(".cache/pip");
759 let _ = std::fs::create_dir_all(&cache_dir);
760
761 let result = fix_pip_cache_ownership(tmp_dir.to_string_lossy().as_ref());
762 assert_eq!(result, PipCacheFixResult::OwnershipCorrect);
764
765 let _ = std::fs::remove_dir_all(&tmp_dir);
767 }
768
769 #[test]
770 fn test_fix_pip_cache_with_temp_dir_no_cache() {
771 let tmp_dir = std::env::temp_dir().join("rusty_stack_test_pip_cache_none");
773 let _ = std::fs::create_dir_all(&tmp_dir);
774
775 let result = fix_pip_cache_ownership(tmp_dir.to_string_lossy().as_ref());
776 assert_eq!(result, PipCacheFixResult::NoCacheDir);
777
778 let _ = std::fs::remove_dir_all(&tmp_dir);
780 }
781
782 #[test]
783 fn test_fix_pip_cache_result_debug() {
784 let result = PipCacheFixResult::Fixed;
786 let debug_str = format!("{:?}", result);
787 assert!(debug_str.contains("Fixed"));
788 }
789}