setup: venv to manage python dependencies automatically

pip-requirements.txt

@@ -0,0 +1,10 @@
+ansiwrap==0.8.4
+capstone==5.0.9
+colorama==0.4.6
+cxxfilt==0.3.0
+Levenshtein==0.27.3
+python-Levenshtein==0.27.3
+RapidFuzz==3.14.5
+textwrap3==0.9.2
+toml==0.10.2
+watchdog==6.0.0

setup_venv.py

@@ -0,0 +1,64 @@
+import os
+import sys
+import subprocess
+from pathlib import Path
+
+USER_FRIENDLY_VENV_PATH = "tools/common/.venv"
+
+def enter_venv():
+    is_already_in_venv = not not os.environ.get("NX_DECOMP_TOOLS_IN_VENV")
+    if is_already_in_venv:
+        return
+    expected_executable = Path(__file__).parent / ".venv" / "bin" / "python"
+    if sys.executable == expected_executable:
+        return
+    setup_python_venv()
+    os.environ["NX_DECOMP_TOOLS_IN_VENV"] = "1"
+    os.execv(expected_executable, [expected_executable, *sys.argv])
+
+def fail(error: str):
+    print(">>> " + error)
+    sys.exit(1)
+
+def setup_python_venv():
+    tools_root = Path(__file__).parent
+    venv_path = tools_root / ".venv"
+    venv_python = venv_path / "bin" / "python"
+    venv_pip = venv_path / "bin" / "pip"
+
+    if not venv_path.is_dir(follow_symlinks=True):
+        if venv_path.exists():
+            fail(f"error: {USER_FRIENDLY_VENV_PATH} exists and is not a directory!")
+        # create venv
+        print(f">>> creating {USER_FRIENDLY_VENV_PATH}")
+        subprocess.check_call([sys.executable, "-m", "venv", venv_path])
+    else:
+        print(f">>> {USER_FRIENDLY_VENV_PATH} is already setup")
+        # still fall through to ensure pip modules are up-to-date
+
+
+    # for some reason just installing with pip install . fail to build
+    # levenshtein, so we will use the shim approach
+    print(">>> installing python dependencies")
+    try:
+        requirements = tools_root / "pip-requirements.txt"
+        subprocess.check_call([
+            venv_pip,
+            "install", 
+            "-r",
+            requirements
+        ])
+        # create shim for asm-differ
+        # note the cd is important so diff_settings is processed correctly
+        asm_differ_shim = venv_path / "bin" / "asm-differ"
+        asm_differ_shim.write_text(f"""#!/usr/bin/env bash
+set -euo pipefail
+PYTHON='{venv_python}'
+cd '{tools_root}'
+exec "$PYTHON" asm-differ/diff.py "$@"
+""")
+        os.chmod(asm_differ_shim, 0o755) # rwxr-xr-x
+
+    except:
+        print(sys.exc_info()[0])
+        fail(f"error: delete {USER_FRIENDLY_VENV_PATH} and try again")

viking/src/tools/check.rs

@@ -9,13 +9,14 @@ use itertools::Itertools;
 use lexopt::prelude::*;
 use rayon::prelude::*;
 use std::cell::RefCell;
-use std::collections::HashMap;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::env;
 use std::fs::File;
 use std::io::{BufRead, BufReader};
+use std::path::PathBuf;
 use std::sync::atomic;
 use std::sync::Mutex;
+use std::process::Command;
 use viking::checks::FunctionChecker;
 use viking::checks::Mismatch;
 use viking::elf;
@@ -174,7 +175,7 @@ All further arguments are forwarded onto asm-differ.
 asm-differ arguments:"
 );
 
-    let differ_path = repo::get_tools_path()?.join("asm-differ").join("diff.py");
+    let differ_path = get_asm_differ_path()?;
 
     // By default, invoking asm-differ using std::process:Process doesn't seem to allow argparse
     // (the python module asm-differ uses to print its help text) to correctly determine the number of columns in the host terminal.
@@ -184,7 +185,7 @@ asm-differ arguments:"
         Err(_) => 240,
     };
 
-    let output = std::process::Command::new(&differ_path)
+    let output = Command::new(&differ_path)
         .current_dir(repo::get_tools_path()?)
         .arg("--help")
         .env("COLUMNS", num_columns.to_string())
@@ -468,14 +469,13 @@ fn check_all(checker: &FunctionChecker, functions: &[functions::Info], args: &Ar
         &new_function_statuses.lock().unwrap(),
         args.version.as_deref(),
     )
-    .with_context(|| "failed to update function statuses")?;
+    .context("failed to update function statuses")?;
 
     if failed.load(atomic::Ordering::Relaxed) {
         bail!("found at least one error");
-    } else {
-        eprintln!("{}", "OK".green().bold());
-        Ok(())
     }
+    eprintln!("{}", "OK".green().bold());
+    Ok(())
 }
 
 #[cold]
@@ -606,8 +606,8 @@ fn show_asm_differ(
     differ_args: &[String],
     version: Option<&str>,
 ) -> Result<()> {
-    let differ_path = repo::get_tools_path()?.join("asm-differ").join("diff.py");
-    let mut cmd = std::process::Command::new(&differ_path);
+    let differ_path = get_asm_differ_path()?;
+    let mut cmd = Command::new(&differ_path);
 
     cmd.current_dir(repo::get_tools_path()?)
         .arg("-I")
@@ -627,6 +627,23 @@ fn show_asm_differ(
     Ok(())
 }
 
+fn get_asm_differ_path() -> Result<PathBuf> {
+    let base_path = repo::get_tools_path()?;
+    let differ_path_venv = {
+        let mut p = base_path.clone();
+        p.extend([".venv", "bin", "asm-differ"]);
+        p
+    };
+    // use the virtual env if one is setup with setup_python_venv()
+    if differ_path_venv.exists() {
+        return Ok(differ_path_venv)
+    }
+    // fallback to the .py entry point
+    let mut p = base_path;
+    p.extend(["asm-differ", "diff.py"]);
+    Ok(p)
+}
+
 fn rediff_function_after_differ(
     functions: &[functions::Info],
     orig_fn: &elf::Function,