use std::fmt;
use std::io::IsTerminal;
use std::path::Path;
use std::process::Command;
#[derive(Debug, Clone, Copy)]
pub struct Colors {
pub red: &'static str,
pub green: &'static str,
pub yellow: &'static str,
pub blue: &'static str,
pub magenta: &'static str,
pub cyan: &'static str,
pub bold: &'static str,
pub reset: &'static str,
}
const COLOR_ENABLED: Colors = Colors {
red: "\x1b[0;31m",
green: "\x1b[0;32m",
yellow: "\x1b[0;33m",
blue: "\x1b[0;34m",
magenta: "\x1b[0;35m",
cyan: "\x1b[0;36m",
bold: "\x1b[1m",
reset: "\x1b[0m",
};
const COLOR_DISABLED: Colors = Colors {
red: "",
green: "",
yellow: "",
blue: "",
magenta: "",
cyan: "",
bold: "",
reset: "",
};
pub fn color_enabled() -> bool {
if std::env::var("NO_COLOR").is_ok() {
return false;
}
std::io::stderr().is_terminal()
}
pub fn get_colors() -> &'static Colors {
if color_enabled() {
&COLOR_ENABLED
} else {
&COLOR_DISABLED
}
}
pub fn command_exists(cmd: &str) -> bool {
Command::new("which")
.arg(cmd)
.output()
.map(|o| o.status.success())
.unwrap_or(false)
}
pub fn resolve_rocminfo_path() -> String {
let standard = Path::new("/opt/rocm/bin/rocminfo");
if standard.exists() {
return standard.to_string_lossy().to_string();
}
if let Ok(entries) = std::fs::read_dir("/opt") {
let mut rocm_dirs: Vec<_> = entries
.filter_map(|e| e.ok())
.filter(|e| e.file_name().to_string_lossy().starts_with("rocm-"))
.collect();
rocm_dirs.sort_by_key(|b| std::cmp::Reverse(b.file_name()));
for dir in rocm_dirs {
let path = dir.path().join("bin/rocminfo");
if path.exists() {
return path.to_string_lossy().to_string();
}
}
}
let arch_path = Path::new("/usr/lib/rocm/bin/rocminfo");
if arch_path.exists() {
return arch_path.to_string_lossy().to_string();
}
"rocminfo".to_string()
}
pub fn resolve_rocprofiler_tool_path() -> Option<String> {
let candidates = [
"/opt/rocm/lib/rocprofiler-sdk/librocprofiler-sdk-tool.so",
"/opt/rocm/lib/librocprofiler-sdk-tool.so",
"/usr/lib/rocm/lib/rocprofiler-sdk/librocprofiler-sdk-tool.so",
"/usr/lib/rocm/lib/librocprofiler-sdk-tool.so",
];
for candidate in &candidates {
if Path::new(candidate).exists() {
return Some(candidate.to_string());
}
}
if let Ok(entries) = std::fs::read_dir("/opt") {
let mut rocm_dirs: Vec<_> = entries
.filter_map(|e| e.ok())
.filter(|e| e.file_name().to_string_lossy().starts_with("rocm-"))
.collect();
rocm_dirs.sort_by_key(|b| std::cmp::Reverse(b.file_name()));
for dir in rocm_dirs {
let base = dir.path();
for subpath in [
"lib/rocprofiler-sdk/librocprofiler-sdk-tool.so",
"lib/librocprofiler-sdk-tool.so",
] {
let path = base.join(subpath);
if path.exists() {
return Some(path.to_string_lossy().to_string());
}
}
}
}
None
}
pub fn print_header(title: &str) {
let c = get_colors();
let width = 57;
let inner = width - 2; let padded = format!("=== {title} ===");
let centered = center_str(&padded, inner);
let border_line = "═".repeat(inner);
let empty_line = " ".repeat(inner);
eprintln!();
eprintln!("{}╔{border_line}╗{}", c.cyan, c.reset);
eprintln!("{}║{empty_line}║{}", c.cyan, c.reset);
eprintln!("{}║{centered}║{}", c.cyan, c.reset);
eprintln!("{}║{empty_line}║{}", c.cyan, c.reset);
eprintln!("{}╚{border_line}╝{}", c.cyan, c.reset);
eprintln!();
}
pub fn print_section(title: &str) {
let c = get_colors();
eprintln!();
eprintln!(
"{}┌─────────────────────────────────────────────────────────┐{}",
c.blue, c.reset
);
eprintln!("{}│ {}{}{}", c.blue, c.bold, title, c.reset);
eprintln!(
"{}└─────────────────────────────────────────────────────────┘{}",
c.blue, c.reset
);
}
pub fn print_step(message: &str) {
let c = get_colors();
eprintln!("{}➤ {}{}", c.magenta, message, c.reset);
}
pub fn print_success(message: &str) {
let c = get_colors();
eprintln!("{}✓ {}{}", c.green, message, c.reset);
}
pub fn print_warning(message: &str) {
let c = get_colors();
eprintln!("{}⚠ {}{}", c.yellow, message, c.reset);
}
pub fn print_error(message: &str) {
let c = get_colors();
eprintln!("{}✗ {}{}", c.red, message, c.reset);
}
pub fn execute_command(cmd: &str, description: &str, dry_run: bool) -> Result<(), String> {
let c = get_colors();
if dry_run {
eprintln!("{}[DRY-RUN]{} {}: ", c.yellow, c.reset, description);
eprintln!(" {}{}{}", c.bold, cmd, c.reset);
return Ok(());
}
print_step(&format!("{description}..."));
let result = Command::new("bash").arg("-c").arg(cmd).output();
match result {
Ok(output) if output.status.success() => {
print_success(&format!("Done: {description}"));
Ok(())
}
Ok(output) => {
let stderr = String::from_utf8_lossy(&output.stderr);
print_error(&format!("Failed: {description}"));
if !stderr.is_empty() {
eprintln!(" {stderr}");
}
Err(format!("Command failed: {cmd}"))
}
Err(e) => {
print_error(&format!("Failed: {description}: {e}"));
Err(format!("Command execution error: {e}"))
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PythonPkgManager {
Uv,
Pip,
}
impl fmt::Display for PythonPkgManager {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
PythonPkgManager::Uv => write!(f, "uv"),
PythonPkgManager::Pip => write!(f, "pip"),
}
}
}
pub fn detect_python_pkg_manager() -> PythonPkgManager {
if command_exists("uv") {
PythonPkgManager::Uv
} else {
PythonPkgManager::Pip
}
}
pub fn pip_install_prefix(python_bin: &str) -> Vec<String> {
let mgr = detect_python_pkg_manager();
match mgr {
PythonPkgManager::Uv => {
let mut args = vec!["uv".to_string(), "pip".to_string(), "install".to_string()];
if std::env::var("UV_PYTHON").is_err() {
args.push("--python".to_string());
args.push(python_bin.to_string());
}
args
}
PythonPkgManager::Pip => {
vec![
python_bin.to_string(),
"-m".to_string(),
"pip".to_string(),
"install".to_string(),
]
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PipCacheFixResult {
NoCacheDir,
OwnershipCorrect,
Fixed,
FixFailed(String),
}
impl std::fmt::Display for PipCacheFixResult {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
PipCacheFixResult::NoCacheDir => write!(f, "No pip cache directory found"),
PipCacheFixResult::OwnershipCorrect => write!(f, "Pip cache ownership already correct"),
PipCacheFixResult::Fixed => write!(f, "Fixed pip cache ownership"),
PipCacheFixResult::FixFailed(e) => write!(f, "Failed to fix pip cache ownership: {e}"),
}
}
}
#[cfg(unix)]
pub fn fix_pip_cache_ownership(user_home: &str) -> PipCacheFixResult {
use std::os::unix::fs::MetadataExt;
let cache_dir = Path::new(user_home).join(".cache/pip");
if !cache_dir.exists() {
return PipCacheFixResult::NoCacheDir;
}
let home_meta = match std::fs::metadata(user_home) {
Ok(m) => m,
Err(_) => return PipCacheFixResult::FixFailed("Cannot stat home directory".into()),
};
let target_uid = home_meta.uid();
let target_gid = home_meta.gid();
let cache_meta = match std::fs::metadata(&cache_dir) {
Ok(m) => m,
Err(_) => return PipCacheFixResult::FixFailed("Cannot stat cache directory".into()),
};
if cache_meta.uid() == target_uid && cache_meta.gid() == target_gid {
return PipCacheFixResult::OwnershipCorrect;
}
let chown_cmd = format!(
"chown -R {}:{} \"{}\"",
target_uid,
target_gid,
cache_dir.display()
);
match Command::new("bash").arg("-c").arg(&chown_cmd).output() {
Ok(output) if output.status.success() => PipCacheFixResult::Fixed,
Ok(output) => {
let stderr = String::from_utf8_lossy(&output.stderr);
PipCacheFixResult::FixFailed(format!("chown failed: {stderr}"))
}
Err(e) => PipCacheFixResult::FixFailed(format!("chown execution error: {e}")),
}
}
#[cfg(not(unix))]
pub fn fix_pip_cache_ownership(_user_home: &str) -> PipCacheFixResult {
PipCacheFixResult::OwnershipCorrect
}
fn center_str(s: &str, width: usize) -> String {
if s.len() >= width {
return s.to_string();
}
let left = (width - s.len()) / 2;
let right = width - s.len() - left;
format!("{}{}{}", " ".repeat(left), s, " ".repeat(right))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_colors_enabled_has_escapes() {
let c = &COLOR_ENABLED;
assert!(!c.red.is_empty());
assert!(!c.green.is_empty());
assert!(!c.reset.is_empty());
assert!(c.red.starts_with('\x1b'));
}
#[test]
fn test_colors_disabled_is_empty() {
let c = &COLOR_DISABLED;
assert!(c.red.is_empty());
assert!(c.green.is_empty());
assert!(c.yellow.is_empty());
assert!(c.blue.is_empty());
assert!(c.magenta.is_empty());
assert!(c.cyan.is_empty());
assert!(c.bold.is_empty());
assert!(c.reset.is_empty());
}
#[test]
fn test_get_colors_returns_valid_ref() {
let _colors = get_colors();
}
#[test]
fn test_command_exists_true() {
assert!(command_exists("bash"));
}
#[test]
fn test_command_exists_false() {
assert!(!command_exists("nonexistent_command_xyz_12345"));
}
#[test]
fn test_command_exists_which() {
assert!(command_exists("which"));
}
#[test]
fn test_print_header_no_panic() {
print_header("Test Title");
}
#[test]
fn test_print_section_no_panic() {
print_section("Test Section");
}
#[test]
fn test_print_step_no_panic() {
print_step("Test step");
}
#[test]
fn test_print_success_no_panic() {
print_success("Test success");
}
#[test]
fn test_print_warning_no_panic() {
print_warning("Test warning");
}
#[test]
fn test_print_error_no_panic() {
print_error("Test error");
}
#[test]
fn test_execute_command_dry_run() {
let result = execute_command("echo hello", "test dry run", true);
assert!(result.is_ok());
}
#[test]
fn test_execute_command_success() {
let result = execute_command("true", "test success", false);
assert!(result.is_ok());
}
#[test]
fn test_execute_command_failure() {
let result = execute_command("false", "test failure", false);
assert!(result.is_err());
}
#[test]
fn test_detect_python_pkg_manager() {
let mgr = detect_python_pkg_manager();
assert!(mgr == PythonPkgManager::Uv || mgr == PythonPkgManager::Pip);
}
#[test]
fn test_python_pkg_manager_display() {
assert_eq!(format!("{}", PythonPkgManager::Uv), "uv");
assert_eq!(format!("{}", PythonPkgManager::Pip), "pip");
}
#[test]
fn test_pip_install_prefix_returns_nonempty() {
let prefix = pip_install_prefix("python3");
assert!(!prefix.is_empty());
assert!(prefix[0] == "uv" || prefix[0] == "python3");
}
#[test]
fn test_center_str_short() {
let result = center_str("hi", 10);
assert_eq!(result.len(), 10);
assert_eq!(result.trim(), "hi");
}
#[test]
fn test_center_str_exact() {
let result = center_str("hello", 5);
assert_eq!(result, "hello");
}
#[test]
fn test_center_str_longer() {
let result = center_str("hello world", 5);
assert_eq!(result, "hello world");
}
#[test]
fn test_resolve_rocminfo_path_returns_string() {
let path = resolve_rocminfo_path();
assert!(
!path.is_empty(),
"resolve_rocminfo_path must return non-empty string"
);
}
#[test]
fn test_resolve_rocminfo_path_prefers_standard_path() {
let path = resolve_rocminfo_path();
if Path::new("/opt/rocm/bin/rocminfo").exists() {
assert_eq!(path, "/opt/rocm/bin/rocminfo");
}
assert!(!path.is_empty());
}
#[test]
fn test_resolve_rocminfo_path_fallback_is_rocminfo() {
let path = resolve_rocminfo_path();
assert!(
path == "rocminfo" || path.contains("rocminfo"),
"Path should contain 'rocminfo', got: {path}"
);
}
#[test]
fn test_pip_cache_fix_result_display_no_cache_dir() {
assert_eq!(
format!("{}", PipCacheFixResult::NoCacheDir),
"No pip cache directory found"
);
}
#[test]
fn test_pip_cache_fix_result_display_ownership_correct() {
assert_eq!(
format!("{}", PipCacheFixResult::OwnershipCorrect),
"Pip cache ownership already correct"
);
}
#[test]
fn test_pip_cache_fix_result_display_fixed() {
assert_eq!(
format!("{}", PipCacheFixResult::Fixed),
"Fixed pip cache ownership"
);
}
#[test]
fn test_pip_cache_fix_result_display_fix_failed() {
let result = PipCacheFixResult::FixFailed("some error".into());
assert!(format!("{}", result).contains("some error"));
}
#[test]
fn test_pip_cache_fix_result_equality() {
assert_eq!(PipCacheFixResult::NoCacheDir, PipCacheFixResult::NoCacheDir);
assert_eq!(
PipCacheFixResult::OwnershipCorrect,
PipCacheFixResult::OwnershipCorrect
);
assert_eq!(PipCacheFixResult::Fixed, PipCacheFixResult::Fixed);
assert_ne!(PipCacheFixResult::NoCacheDir, PipCacheFixResult::Fixed);
}
#[test]
fn test_fix_pip_cache_no_cache_dir() {
let result = fix_pip_cache_ownership("/tmp/nonexistent_test_dir_xyz");
assert_eq!(result, PipCacheFixResult::NoCacheDir);
}
#[test]
fn test_fix_pip_cache_ownership_correct() {
let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".into());
let result = fix_pip_cache_ownership(&home);
assert!(
result == PipCacheFixResult::NoCacheDir
|| result == PipCacheFixResult::OwnershipCorrect
|| result == PipCacheFixResult::Fixed,
"Expected NoCacheDir, OwnershipCorrect, or Fixed, got: {result}"
);
}
#[test]
fn test_fix_pip_cache_with_temp_dir_correct_ownership() {
let tmp_dir = std::env::temp_dir().join("rusty_stack_test_pip_cache_correct");
let cache_dir = tmp_dir.join(".cache/pip");
let _ = std::fs::create_dir_all(&cache_dir);
let result = fix_pip_cache_ownership(tmp_dir.to_string_lossy().as_ref());
assert_eq!(result, PipCacheFixResult::OwnershipCorrect);
let _ = std::fs::remove_dir_all(&tmp_dir);
}
#[test]
fn test_fix_pip_cache_with_temp_dir_no_cache() {
let tmp_dir = std::env::temp_dir().join("rusty_stack_test_pip_cache_none");
let _ = std::fs::create_dir_all(&tmp_dir);
let result = fix_pip_cache_ownership(tmp_dir.to_string_lossy().as_ref());
assert_eq!(result, PipCacheFixResult::NoCacheDir);
let _ = std::fs::remove_dir_all(&tmp_dir);
}
#[test]
fn test_fix_pip_cache_result_debug() {
let result = PipCacheFixResult::Fixed;
let debug_str = format!("{:?}", result);
assert!(debug_str.contains("Fixed"));
}
}