use std::env::consts::EXE_EXTENSION;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::{Arc, LazyLock};
use anyhow::{Context, Result};
use mea::once::OnceMap;
use prek_consts::env_vars::EnvVars;
use prek_consts::prepend_paths;
use rustc_hash::FxBuildHasher;
use serde::Deserialize;
use tracing::{debug, trace};
use crate::cli::reporter::{HookInstallReporter, HookRunReporter};
use crate::hook::InstalledHook;
use crate::hook::{Hook, InstallInfo};
use crate::languages::LanguageImpl;
use crate::languages::python::PythonRequest;
use crate::languages::python::uv::Uv;
use crate::languages::version::LanguageRequest;
use crate::process;
use crate::process::Cmd;
use crate::run::run_by_batch;
use crate::store::{Store, ToolBucket};
#[derive(Debug, Copy, Clone)]
pub(crate) struct Python;
pub(crate) struct PythonInfo {
pub(crate) version: semver::Version,
pub(crate) python_exec: PathBuf,
}
#[derive(Debug, Clone, thiserror::Error)]
pub(crate) enum PythonInfoError {
#[error("Failed to parse Python info JSON: {0}")]
Parse(String),
#[error("Failed to query Python info: {0}")]
Query(String),
#[error("{0}")]
Message(String),
}
static PYTHON_INFO_CACHE: LazyLock<OnceMap<PathBuf, Arc<PythonInfo>, FxBuildHasher>> =
LazyLock::new(|| OnceMap::with_hasher(FxBuildHasher));
async fn query_python_info(python: &Path) -> Result<PythonInfo, PythonInfoError> {
#[derive(Deserialize)]
struct QueryPythonInfo {
version: semver::Version,
base_exec_prefix: PathBuf,
}
static QUERY_PYTHON_INFO: &str = indoc::indoc! {r#"
import sys, json
info = {
"version": ".".join(map(str, sys.version_info[:3])),
"base_exec_prefix": sys.base_exec_prefix,
}
print(json.dumps(info))
"#};
let stdout = Cmd::new(python, "python -c")
.arg("-I")
.arg("-c")
.arg(QUERY_PYTHON_INFO)
.check(true)
.output()
.await
.map_err(|err| PythonInfoError::Query(err.to_string()))?
.stdout;
let info: QueryPythonInfo =
serde_json::from_slice(&stdout).map_err(|err| PythonInfoError::Parse(err.to_string()))?;
let python_exec = python_exec(&info.base_exec_prefix);
Ok(PythonInfo {
version: info.version,
python_exec,
})
}
pub(crate) async fn query_python_info_cached(
python: &Path,
) -> Result<Arc<PythonInfo>, PythonInfoError> {
let python = fs::canonicalize(python).unwrap_or_else(|_| python.to_path_buf());
PYTHON_INFO_CACHE
.try_compute(python.clone(), async move || {
let info = query_python_info(&python).await?;
Ok(Arc::new(info))
})
.await
}
impl LanguageImpl for Python {
async fn install(
&self,
hook: Arc<Hook>,
store: &Store,
reporter: &HookInstallReporter,
) -> Result<InstalledHook> {
let progress = reporter.on_install_start(&hook);
let uv_dir = store.tools_path(ToolBucket::Uv);
let uv = Uv::install(store, &uv_dir)
.await
.context("Failed to install uv")?;
let mut info = InstallInfo::new(
hook.language,
hook.env_key_dependencies().clone(),
&store.hooks_dir(),
)?;
debug!(%hook, target = %info.env_path.display(), "Installing environment");
Self::create_venv(&uv, store, &info, &hook.language_request)
.await
.context("Failed to create Python virtual environment")?;
let mut pip_install = Self::pip_install_command(&uv, store, &info.env_path);
if let Some(repo_path) = hook.repo_path() {
trace!(
"Installing dependencies from repo path: {}",
repo_path.display()
);
pip_install
.arg("--directory")
.arg(repo_path)
.arg(".")
.args(&hook.additional_dependencies)
.output()
.await?;
} else if !hook.additional_dependencies.is_empty() {
trace!(
"Installing additional dependencies: {:?}",
hook.additional_dependencies
);
pip_install
.args(&hook.additional_dependencies)
.output()
.await?;
} else {
debug!("No dependencies to install");
}
let python = python_exec(&info.env_path);
let python_info = query_python_info(&python)
.await
.context("Failed to query Python info")?;
info.with_language_version(python_info.version)
.with_toolchain(python_info.python_exec);
info.persist_env_path();
reporter.on_install_complete(progress);
Ok(InstalledHook::Installed {
hook,
info: Arc::new(info),
})
}
async fn check_health(&self, info: &InstallInfo) -> Result<()> {
let python = python_exec(&info.env_path);
let python_info = query_python_info_cached(&python)
.await
.context("Failed to query Python info")?;
if python_info.version != info.language_version {
anyhow::bail!(
"Python version mismatch: expected {}, found {}",
info.language_version,
python_info.version
);
}
Ok(())
}
async fn run(
&self,
hook: &InstalledHook,
filenames: &[&Path],
store: &Store,
reporter: &HookRunReporter,
) -> Result<(i32, Vec<u8>)> {
let progress = reporter.on_run_start(hook, filenames.len());
let env_dir = hook.env_path().expect("Python must have env path");
let new_path = prepend_paths(&[&bin_dir(env_dir)]).context("Failed to join PATH")?;
let entry = hook.entry.resolve(Some(&new_path), store)?;
let run = async |batch: &[&Path]| {
let mut output = Cmd::new(&entry[0], "python hook")
.current_dir(hook.work_dir())
.args(&entry[1..])
.env(EnvVars::VIRTUAL_ENV, env_dir)
.env(EnvVars::PATH, &new_path)
.env_remove(EnvVars::PYTHONHOME)
.envs(&hook.env)
.args(&hook.args)
.args(batch)
.check(false)
.stdin(Stdio::null())
.pty_output()
.await?;
reporter.on_run_progress(progress, batch.len() as u64);
output.stdout.extend(output.stderr);
let code = output.status.code().unwrap_or(1);
anyhow::Ok((code, output.stdout))
};
let results = run_by_batch(hook, filenames, entry.argv(), run).await?;
reporter.on_run_complete(progress);
let mut combined_status = 0;
let mut combined_output = Vec::new();
for (code, output) in results {
combined_status |= code;
combined_output.extend(output);
}
Ok((combined_status, combined_output))
}
}
fn to_uv_python_request(request: &LanguageRequest) -> Option<String> {
match request {
LanguageRequest::Any { .. } => None,
LanguageRequest::Python(request) => match request {
PythonRequest::Any => None,
PythonRequest::Major(major) => Some(format!("{major}")),
PythonRequest::MajorMinor(major, minor) => Some(format!("{major}.{minor}")),
PythonRequest::MajorMinorPatch(major, minor, patch) => {
Some(format!("{major}.{minor}.{patch}"))
}
PythonRequest::Range(_, raw) => Some(raw.clone()),
PythonRequest::Path(path) => Some(path.to_string_lossy().to_string()),
},
_ => unreachable!(),
}
}
impl Python {
fn remove_uv_python_override_envs(cmd: &mut Cmd) -> &mut Cmd {
cmd.env_remove(EnvVars::UV_PYTHON)
.env_remove(EnvVars::UV_SYSTEM_PYTHON)
.env_remove(EnvVars::UV_MANAGED_PYTHON)
.env_remove(EnvVars::UV_NO_MANAGED_PYTHON)
}
fn pip_install_command(uv: &Uv, store: &Store, env_path: &Path) -> Cmd {
let mut cmd = uv.cmd("uv pip", store);
cmd.arg("pip")
.arg("install")
.args(["--project", "/"])
.env(EnvVars::VIRTUAL_ENV, env_path);
Self::remove_uv_python_override_envs(&mut cmd)
.remove_git_envs()
.check(true);
cmd
}
async fn create_venv(
uv: &Uv,
store: &Store,
info: &InstallInfo,
python_request: &LanguageRequest,
) -> Result<()> {
match Self::create_venv_command(uv, store, info, python_request, false, false)
.check(true)
.output()
.await
{
Ok(_) => {
debug!(
"Venv created successfully with no downloads: `{}`",
info.env_path.display()
);
Ok(())
}
Err(e @ process::Error::Status { .. }) => {
if Self::can_retry_with_downloads(&e) {
if !python_request.allows_download() {
anyhow::bail!(
"No suitable system Python version found and downloads are disabled"
);
}
debug!(
"Retrying venv creation with managed Python downloads: `{}`",
info.env_path.display()
);
Self::create_venv_command(uv, store, info, python_request, true, true)
.check(true)
.output()
.await?;
return Ok(());
}
Err(e.into())
}
Err(e) => {
debug!("Failed to create venv `{}`: {e}", info.env_path.display());
Err(e.into())
}
}
}
fn create_venv_command(
uv: &Uv,
store: &Store,
info: &InstallInfo,
python_request: &LanguageRequest,
set_install_dir: bool,
allow_downloads: bool,
) -> Cmd {
let mut cmd = uv.cmd("create venv", store);
cmd.arg("venv")
.arg(&info.env_path)
.args(["--python-preference", "managed"])
.arg("--no-project")
.args(["--project", "/"]);
Self::remove_uv_python_override_envs(&mut cmd);
if set_install_dir {
cmd.env(
EnvVars::UV_PYTHON_INSTALL_DIR,
store.tools_path(ToolBucket::Python),
);
}
if allow_downloads {
cmd.arg("--allow-python-downloads");
} else {
cmd.arg("--no-python-downloads");
}
if let Some(python) = to_uv_python_request(python_request) {
cmd.arg("--python").arg(python);
}
cmd
}
fn can_retry_with_downloads(error: &process::Error) -> bool {
let process::Error::Status {
error:
process::StatusError {
output: Some(output),
..
},
..
} = error
else {
return false;
};
let stderr = String::from_utf8_lossy(&output.stderr);
stderr.contains("A managed Python download is available")
}
}
fn bin_dir(venv: &Path) -> PathBuf {
if cfg!(windows) {
venv.join("Scripts")
} else {
venv.join("bin")
}
}
pub(crate) fn python_exec(venv: &Path) -> PathBuf {
bin_dir(venv).join("python").with_extension(EXE_EXTENSION)
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::path::PathBuf;
use prek_consts::env_vars::EnvVars;
use rustc_hash::FxHashSet;
use super::Python;
use crate::config::Language;
use crate::hook::InstallInfo;
use crate::languages::python::uv::Uv;
use crate::languages::version::LanguageRequest;
use crate::store::Store;
fn setup_test_install() -> (tempfile::TempDir, Uv, Store, InstallInfo) {
let temp = tempfile::tempdir().expect("create tempdir");
let hooks_dir = temp.path().join("hooks");
fs_err::create_dir_all(&hooks_dir).expect("create hooks dir");
let info = InstallInfo::new(Language::Python, FxHashSet::default(), &hooks_dir)
.expect("create install info");
let store = Store::from_path(temp.path().join("store"));
let uv = Uv::new(PathBuf::from("uv"));
(temp, uv, store, info)
}
fn env_map(cmd: &crate::process::Cmd) -> HashMap<String, Option<String>> {
cmd.get_envs()
.map(|(key, val)| {
(
key.to_string_lossy().into_owned(),
val.map(|v| v.to_string_lossy().into_owned()),
)
})
.collect()
}
#[test]
fn create_venv_command_removes_uv_system_python_override() {
let (_temp, uv, store, info) = setup_test_install();
let request = LanguageRequest::Any { system_only: false };
let cmd = Python::create_venv_command(&uv, &store, &info, &request, false, false);
let envs = env_map(&cmd);
assert_eq!(envs.get(EnvVars::UV_SYSTEM_PYTHON), Some(&None));
assert_eq!(envs.get(EnvVars::UV_PYTHON), Some(&None));
assert_eq!(envs.get(EnvVars::UV_MANAGED_PYTHON), Some(&None));
assert_eq!(envs.get(EnvVars::UV_NO_MANAGED_PYTHON), Some(&None));
}
#[test]
fn pip_install_command_removes_uv_system_python_override() {
let (_temp, uv, store, info) = setup_test_install();
let cmd = Python::pip_install_command(&uv, &store, &info.env_path);
let envs = env_map(&cmd);
assert_eq!(envs.get(EnvVars::UV_SYSTEM_PYTHON), Some(&None));
assert_eq!(envs.get(EnvVars::UV_PYTHON), Some(&None));
assert_eq!(envs.get(EnvVars::UV_MANAGED_PYTHON), Some(&None));
assert_eq!(envs.get(EnvVars::UV_NO_MANAGED_PYTHON), Some(&None));
}
}