Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions crates/uv-fs/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,36 @@ pub fn replace_symlink(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> std::io:
}
}

/// Create a symlink at `dst` pointing to `src`.
///
/// On Windows, this uses the `junction` crate to create a junction point.
///
/// Note that because junctions are used, the source must be a directory.
#[cfg(windows)]
pub fn create_symlink(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> std::io::Result<()> {
// If the source is a file, we can't create a junction
if src.as_ref().is_file() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"Cannot create a junction for {}: is not a directory",
src.as_ref().display()
),
));
}

junction::create(
dunce::simplified(src.as_ref()),
dunce::simplified(dst.as_ref()),
)
}

/// Create a symlink at `dst` pointing to `src`.
#[cfg(unix)]
pub fn create_symlink(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> std::io::Result<()> {
fs_err::os::unix::fs::symlink(src.as_ref(), dst.as_ref())
}

#[cfg(unix)]
pub fn remove_symlink(path: impl AsRef<Path>) -> std::io::Result<()> {
fs_err::remove_file(path.as_ref())
Expand Down
61 changes: 61 additions & 0 deletions crates/uv-trampoline-builder/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ const MAGIC_NUMBER_SIZE: usize = 4;
pub struct Launcher {
pub kind: LauncherKind,
pub python_path: PathBuf,
payload: Vec<u8>,
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In order to allow rewriting the launcher, we need to stash the payload. I think this is generally pretty small. If we are worried about the overhead, we could store the path to the original launcher and load it again when needed.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you want to get fancy you could probably store Box<[u8]> here. Seems totally not necessary though.

}

impl Launcher {
Expand Down Expand Up @@ -109,11 +110,69 @@ impl Launcher {
String::from_utf8(buffer).map_err(|err| Error::InvalidPath(err.utf8_error()))?,
);

#[allow(clippy::cast_possible_truncation)]
let file_size = {
let raw_length = file
.seek(io::SeekFrom::End(0))
.map_err(|e| Error::InvalidLauncherSeek("size probe".into(), 0, e))?;

if raw_length > usize::MAX as u64 {
return Err(Error::InvalidDataLength(raw_length));
}

// SAFETY: Above we guarantee the length is less than uszie
raw_length as usize
};

// Read the payload
file.seek(io::SeekFrom::Start(0))
.map_err(|e| Error::InvalidLauncherSeek("rewind".into(), 0, e))?;
let payload_len =
file_size.saturating_sub(MAGIC_NUMBER_SIZE + PATH_LENGTH_SIZE + path_length);
let mut buffer = vec![0u8; payload_len];
file.read_exact(&mut buffer)
.map_err(|err| Error::InvalidLauncherRead("payload".into(), err))?;

Ok(Some(Self {
kind,
payload: buffer,
python_path: path,
}))
}

pub fn to_file(self, path: &Path) -> Result<(), Error> {
let python_path = self.python_path.simplified_display().to_string();

if python_path.len() > MAX_PATH_LENGTH as usize {
return Err(Error::InvalidPathLength(
u32::try_from(python_path.len()).expect("path length already checked"),
));
}

let mut launcher: Vec<u8> = Vec::with_capacity(
self.payload.len() + python_path.len() + PATH_LENGTH_SIZE + MAGIC_NUMBER_SIZE,
);
launcher.extend_from_slice(&self.payload);
launcher.extend_from_slice(python_path.as_bytes());
launcher.extend_from_slice(
&u32::try_from(python_path.len())
.expect("file path should be smaller than 4GB")
.to_le_bytes(),
);
launcher.extend_from_slice(self.kind.magic_number());

fs_err::write(path, launcher)?;
Ok(())
}

#[must_use]
pub fn with_python_path(self, path: PathBuf) -> Self {
Self {
kind: self.kind,
payload: self.payload,
python_path: path,
}
}
}

/// The kind of trampoline launcher to create.
Expand Down Expand Up @@ -177,6 +236,8 @@ pub enum Error {
Io(#[from] io::Error),
#[error("Only paths with a length up to 32KB are supported but found a length of {0} bytes")]
InvalidPathLength(u32),
#[error("Only data with a length up to usize is supported but found a length of {0} bytes")]
InvalidDataLength(u64),
#[error("Failed to parse executable path")]
InvalidPath(#[source] Utf8Error),
#[error("Failed to seek to {0} at offset {1}")]
Expand Down
14 changes: 14 additions & 0 deletions crates/uv/src/commands/project/environment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,20 @@ impl EphemeralEnvironment {
)?;
Ok(())
}

/// Returns the path to the environment's scripts directory.
pub(crate) fn scripts(&self) -> &Path {
self.0.scripts()
}

/// Returns the path to the environment's Python executable.
pub(crate) fn sys_executable(&self) -> &Path {
self.0.interpreter().sys_executable()
}

pub(crate) fn sys_prefix(&self) -> &Path {
self.0.interpreter().sys_prefix()
}
}

/// A [`PythonEnvironment`] stored in the cache.
Expand Down
141 changes: 139 additions & 2 deletions crates/uv/src/commands/project/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use futures::StreamExt;
use itertools::Itertools;
use owo_colors::OwoColorize;
use tokio::process::Command;
use tracing::{debug, warn};
use tracing::{debug, trace, warn};
use url::Url;

use uv_cache::Cache;
Expand All @@ -22,7 +22,7 @@ use uv_configuration::{
};
use uv_distribution_types::Requirement;
use uv_fs::which::is_executable;
use uv_fs::{PythonExt, Simplified};
use uv_fs::{PythonExt, Simplified, create_symlink, symlink_or_copy_file};
use uv_installer::{SatisfiesResult, SitePackages};
use uv_normalize::{DefaultExtras, DefaultGroups, PackageName};
use uv_python::{
Expand Down Expand Up @@ -1071,6 +1071,71 @@ hint: If you are running a script with `{}` in the shebang, you may need to incl
requirements_site_packages.escape_for_python(),
))?;

// N.B. The order here matters — earlier interpreters take precedence over the
// later ones.
for interpreter in [requirements_env.interpreter(), &base_interpreter] {
// Copy each entrypoint from the base environments to the ephemeral environment.
for entry in fs_err::read_dir(interpreter.scripts())? {
let entry = entry?;

if !entry.file_type()?.is_file() {
continue;
}

copy_entrypoint(
&entry.path(),
&ephemeral_env.scripts().join(entry.file_name()),
interpreter.sys_executable(),
ephemeral_env.sys_executable(),
)?;
}

// Link common data directories from the base environment to the ephemeral
// environment. This is critical for Jupyter Lab, which cannot operate without the
// files it writes to `<prefix>/share`.
//
// We only perform a shallow merge here, so `<prefix>/share/<name>` will only be
// written once and contents beyond the first level, e.g.,
// `<prefix>/share/<name>/foo` will not be merged across parent interpreters.
for dir in &["etc", "share"] {
let entries = match fs_err::read_dir(interpreter.sys_prefix().join(dir)) {
Ok(entries) => entries,
// Skip missing directories
Err(err) if err.kind() == std::io::ErrorKind::NotFound => continue,
Err(err) => return Err(err.into()),
};
fs_err::create_dir_all(ephemeral_env.sys_prefix().join(dir))?;
for entry in entries {
let entry = entry?;
let target = ephemeral_env.sys_prefix().join(dir).join(entry.file_name());

if entry.file_type()?.is_file() {
match symlink_or_copy_file(entry.path(), &target) {
Ok(()) => trace!(
"Created link for {} -> {}",
target.user_display(),
entry.path().user_display()
),
Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {}
Err(err) => return Err(err.into()),
}
} else if entry.file_type()?.is_dir() {
match create_symlink(entry.path(), &target) {
Ok(()) => trace!(
"Created link for {} -> {}",
target.user_display(),
entry.path().user_display()
),
Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {}
Err(err) => return Err(err.into()),
}
} else {
trace!("Skipping link of entry: {}", entry.path().user_display());
}
}
}
}

// Write the `sys.prefix` of the parent environment to the `extends-environment` key of the `pyvenv.cfg`
// file. This helps out static-analysis tools such as ty (see docs on
// `CachedEnvironment::set_parent_environment`).
Expand All @@ -1095,6 +1160,10 @@ hint: If you are running a script with `{}` in the shebang, you may need to incl
// Cast to `PythonEnvironment`.
let ephemeral_env = ephemeral_env.map(PythonEnvironment::from);

if let Some(e) = ephemeral_env.as_ref() {
debug!("Using ephemeral environment at: {}", e.scripts().display());
}

// Determine the Python interpreter to use for the command, if necessary.
let interpreter = ephemeral_env
.as_ref()
Expand Down Expand Up @@ -1669,3 +1738,71 @@ fn read_recursion_depth_from_environment_variable() -> anyhow::Result<u32> {
.parse::<u32>()
.with_context(|| format!("invalid value for {}", EnvVars::UV_RUN_RECURSION_DEPTH))
}

#[cfg(unix)]
fn copy_entrypoint(
source: &Path,
target: &Path,
previous_executable: &Path,
python_executable: &Path,
) -> anyhow::Result<()> {
use std::os::unix::fs::PermissionsExt;

let contents = fs_err::read_to_string(source)?;
let expected = r#"#!/bin/sh
'''exec' "$(dirname -- "$(realpath -- "$0")")"/'python' "$0" "$@"
' '''
"#;

// Only rewrite entrypoints that use the expected shebang.
let Some(contents) = contents
.strip_prefix(expected)
.or_else(|| contents.strip_prefix(&format!("#!{}", previous_executable.display())))
else {
debug!(
"Skipping copy of entrypoint at {}: does not start with expected shebang",
source.user_display()
);
return Ok(());
};

let contents = format!("#!{}\n{}", python_executable.display(), contents);
fs_err::write(target, &contents)?;

match fs_err::write(target, &contents) {
Ok(()) => trace!("Updated entrypoint at {}", target.user_display()),
Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {}
Err(err) => return Err(err.into()),
}

let mut perms = fs_err::metadata(target)?.permissions();
perms.set_mode(0o755);
fs_err::set_permissions(target, perms)?;

Ok(())
}

#[cfg(windows)]
fn copy_entrypoint(
source: &Path,
target: &Path,
_previous_executable: &Path,
python_executable: &Path,
) -> anyhow::Result<()> {
use uv_trampoline_builder::Launcher;

let Some(launcher) = Launcher::try_from_path(source)? else {
return Ok(());
};

let launcher = launcher.with_python_path(python_executable.to_path_buf());

match launcher.to_file(target) {
Ok(()) => trace!("Updated entrypoint at {}", target.user_display()),
Err(uv_trampoline_builder::Error::Io(err))
if err.kind() == std::io::ErrorKind::AlreadyExists => {}
Err(err) => return Err(err.into()),
}

Ok(())
}
Loading
Loading