uv/crates/uv-virtualenv/src/bare.rs

//! Create a bare virtualenv without any packages install

use std::env;
use std::env::consts::EXE_SUFFIX;
use std::io;
use std::io::{BufWriter, Write};

use camino::{FromPathBufError, Utf8Path, Utf8PathBuf};
use fs_err as fs;
use fs_err::File;
use tracing::info;
use uv_fs::Simplified;

use uv_interpreter::{Interpreter, SysconfigPaths, Virtualenv};

use crate::{Error, Prompt};

/// The bash activate scripts with the venv dependent paths patches out
const ACTIVATE_TEMPLATES: &[(&str, &str)] = &[
    ("activate", include_str!("activator/activate")),
    ("activate.csh", include_str!("activator/activate.csh")),
    ("activate.fish", include_str!("activator/activate.fish")),
    ("activate.nu", include_str!("activator/activate.nu")),
    ("activate.ps1", include_str!("activator/activate.ps1")),
    ("activate.bat", include_str!("activator/activate.bat")),
    ("deactivate.bat", include_str!("activator/deactivate.bat")),
    ("pydoc.bat", include_str!("activator/pydoc.bat")),
    (
        "activate_this.py",
        include_str!("activator/activate_this.py"),
    ),
];
const VIRTUALENV_PATCH: &str = include_str!("_virtualenv.py");

/// Very basic `.cfg` file format writer.
fn write_cfg(f: &mut impl Write, data: &[(String, String)]) -> io::Result<()> {
    for (key, value) in data {
        writeln!(f, "{key} = {value}")?;
    }
    Ok(())
}

/// Write all the files that belong to a venv without any packages installed.
pub fn create_bare_venv(
    location: &Utf8Path,
    interpreter: &Interpreter,
    prompt: Prompt,
    system_site_packages: bool,
    extra_cfg: Vec<(String, String)>,
) -> Result<Virtualenv, Error> {
    // We have to canonicalize the interpreter path, otherwise the home is set to the venv dir instead of the real root.
    // This would make python-build-standalone fail with the encodings module not being found because its home is wrong.
    let base_python: Utf8PathBuf = fs_err::canonicalize(interpreter.sys_executable())?
        .try_into()
        .map_err(|err: FromPathBufError| err.into_io_error())?;

    // Validate the existing location.
    match location.metadata() {
        Ok(metadata) => {
            if metadata.is_file() {
                return Err(Error::IO(io::Error::new(
                    io::ErrorKind::AlreadyExists,
                    format!("File exists at `{location}`"),
                )));
            } else if metadata.is_dir() {
                if location.join("pyvenv.cfg").is_file() {
                    info!("Removing existing directory");
                    fs::remove_dir_all(location)?;
                    fs::create_dir_all(location)?;
                } else if location
                    .read_dir()
                    .is_ok_and(|mut dir| dir.next().is_none())
                {
                    info!("Ignoring empty directory");
                } else {
                    return Err(Error::IO(io::Error::new(
                        io::ErrorKind::AlreadyExists,
                        format!("The directory `{location}` exists, but it's not a virtualenv"),
                    )));
                }
            }
        }
        Err(err) if err.kind() == io::ErrorKind::NotFound => {
            fs::create_dir_all(location)?;
        }
        Err(err) => return Err(Error::IO(err)),
    }

    let location = location.canonicalize_utf8()?;

    let bin_name = if cfg!(unix) {
        "bin"
    } else if cfg!(windows) {
        "Scripts"
    } else {
        unimplemented!("Only Windows and Unix are supported")
    };
    let scripts = location.join(bin_name);
    let prompt = match prompt {
        Prompt::CurrentDirectoryName => env::current_dir()?
            .file_name()
            .map(|name| name.to_string_lossy().to_string()),
        Prompt::Static(value) => Some(value),
        Prompt::None => None,
    };

    // Add the CACHEDIR.TAG.
    cachedir::ensure_tag(&location)?;

    // Create a `.gitignore` file to ignore all files in the venv.
    fs::write(location.join(".gitignore"), "*")?;

    // Different names for the python interpreter
    fs::create_dir(&scripts)?;
    let executable = scripts.join(format!("python{EXE_SUFFIX}"));

    // No symlinking on Windows, at least not on a regular non-dev non-admin Windows install.
    #[cfg(unix)]
    {
        use fs_err::os::unix::fs::symlink;

        symlink(&base_python, &executable)?;
        symlink(
            "python",
            scripts.join(format!("python{}", interpreter.python_major())),
        )?;
        symlink(
            "python",
            scripts.join(format!(
                "python{}.{}",
                interpreter.python_major(),
                interpreter.python_minor(),
            )),
        )?;
    }

    #[cfg(windows)]
    {
        // https://github.com/python/cpython/blob/d457345bbc6414db0443819290b04a9a4333313d/Lib/venv/__init__.py#L261-L267
        // https://github.com/pypa/virtualenv/blob/d9fdf48d69f0d0ca56140cf0381edbb5d6fe09f5/src/virtualenv/create/via_global_ref/builtin/cpython/cpython3.py#L78-L83
        // There's two kinds of applications on windows: Those that allocate a console (python.exe) and those that
        // don't because they use window(s) (pythonw.exe).
        for python_exe in ["python.exe", "pythonw.exe"] {
            let shim = interpreter
                .stdlib()
                .join("venv")
                .join("scripts")
                .join("nt")
                .join(python_exe);
            fs_err::copy(shim, scripts.join(python_exe))?;
        }
    }
    #[cfg(not(any(unix, windows)))]
    {
        compile_error!("Only Windows and Unix are supported")
    }

    // Add all the activate scripts for different shells
    for (name, template) in ACTIVATE_TEMPLATES {
        let relative_site_packages = if cfg!(unix) {
            format!(
                "../lib/{}{}.{}/site-packages",
                interpreter.site_packages_python(),
                interpreter.python_major(),
                interpreter.python_minor(),
            )
        } else if cfg!(windows) {
            "../Lib/site-packages".to_string()
        } else {
            unimplemented!("Only Windows and Unix are supported")
        };
        let activator = template
            .replace(
                "{{ VIRTUAL_ENV_DIR }}",
                // SAFETY: `unwrap` is guaranteed to succeed because `location` is an `Utf8PathBuf`.
                location.simplified().to_str().unwrap(),
            )
            .replace("{{ BIN_NAME }}", bin_name)
            .replace(
                "{{ VIRTUAL_PROMPT }}",
                prompt.as_deref().unwrap_or_default(),
            )
            .replace(
                "{{ RELATIVE_SITE_PACKAGES }}",
                relative_site_packages.as_str(),
            );
        fs::write(scripts.join(name), activator)?;
    }

    // pyvenv.cfg
    let python_home = if cfg!(unix) {
        // On Linux and Mac, Python is symlinked so the base home is the parent of the resolved-by-canonicalize path.
        base_python
            .parent()
            .ok_or_else(|| {
                io::Error::new(
                    io::ErrorKind::NotFound,
                    "The python interpreter needs to have a parent directory",
                )
            })?
            .to_string()
    } else if cfg!(windows) {
        // `virtualenv` seems to rely on the undocumented, private `sys._base_executable`. When I tried,
        // `sys.base_prefix` was the same as the parent of `sys._base_executable`, but a much simpler logic and
        // documented.
        // https://github.com/pypa/virtualenv/blob/d9fdf48d69f0d0ca56140cf0381edbb5d6fe09f5/src/virtualenv/discovery/py_info.py#L136-L156
        interpreter.base_prefix().display().to_string()
    } else {
        unimplemented!("Only Windows and Unix are supported")
    };

    // Validate extra_cfg
    let reserved_keys = [
        "home",
        "implementation",
        "version_info",
        "include-system-site-packages",
        "base-prefix",
        "base-exec-prefix",
        "base-executable",
        "prompt",
    ];
    for (key, _) in &extra_cfg {
        if reserved_keys.contains(&key.as_str()) {
            return Err(Error::ReservedConfigKey(key.to_string()));
        }
    }

    let mut pyvenv_cfg_data: Vec<(String, String)> = vec![
        ("home".to_string(), python_home),
        (
            "implementation".to_string(),
            interpreter.markers().platform_python_implementation.clone(),
        ),
        (
            "version_info".to_string(),
            interpreter.markers().python_full_version.string.clone(),
        ),
        (
            "include-system-site-packages".to_string(),
            if system_site_packages {
                "true".to_string()
            } else {
                "false".to_string()
            },
        ),
        (
            "base-prefix".to_string(),
            interpreter.base_prefix().to_string_lossy().to_string(),
        ),
        (
            "base-exec-prefix".to_string(),
            interpreter.base_exec_prefix().to_string_lossy().to_string(),
        ),
        ("base-executable".to_string(), base_python.to_string()),
    ]
    .into_iter()
    .chain(extra_cfg)
    .collect();

    if let Some(prompt) = prompt {
        pyvenv_cfg_data.push(("prompt".to_string(), prompt));
    }

    let mut pyvenv_cfg = BufWriter::new(File::create(location.join("pyvenv.cfg"))?);
    write_cfg(&mut pyvenv_cfg, &pyvenv_cfg_data)?;
    drop(pyvenv_cfg);

    // Construct the path to the `site-packages` directory.
    let site_packages = if cfg!(unix) {
        location
            .join("lib")
            .join(format!(
                "{}{}.{}",
                interpreter.site_packages_python(),
                interpreter.python_major(),
                interpreter.python_minor(),
            ))
            .join("site-packages")
    } else if cfg!(windows) {
        location.join("Lib").join("site-packages")
    } else {
        unimplemented!("Only Windows and Unix are supported")
    };

    // Construct the path to the `platstdlib` directory.
    let platstdlib = if cfg!(windows) {
        location.join("Lib")
    } else {
        location
            .join("lib")
            .join(format!(
                "{}{}.{}",
                interpreter.site_packages_python(),
                interpreter.python_major(),
                interpreter.python_minor()
            ))
            .join("site-packages")
    };

    // Populate `site-packages` with a `_virtualenv.py` file.
    fs::create_dir_all(&site_packages)?;
    fs::write(site_packages.join("_virtualenv.py"), VIRTUALENV_PATCH)?;
    fs::write(site_packages.join("_virtualenv.pth"), "import _virtualenv")?;

    Ok(Virtualenv {
        root: location.clone().into_std_path_buf(),
        executable: executable.into_std_path_buf(),
        sysconfig_paths: SysconfigPaths {
            // Paths that were already constructed above.
            scripts: scripts.into_std_path_buf(),
            platstdlib: platstdlib.into_std_path_buf(),
            // Set `purelib` and `platlib` to the same value.
            purelib: site_packages.clone().into_std_path_buf(),
            platlib: site_packages.into_std_path_buf(),
            // Inherited from the interpreter.
            stdlib: interpreter.stdlib().to_path_buf(),
            include: interpreter.include().to_path_buf(),
            platinclude: interpreter.platinclude().to_path_buf(),
            data: location.into_std_path_buf(),
        },
    })
}