Skip to content
Merged
Show file tree
Hide file tree
Changes from all 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
2 changes: 1 addition & 1 deletion crates/uv-python/src/environment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ impl PythonEnvironment {
/// N.B. This function also works for system Python environments and users depend on this.
pub fn from_root(root: impl AsRef<Path>, cache: &Cache) -> Result<Self, Error> {
debug!(
"Checking for Python environment at `{}`",
"Checking for Python environment at: `{}`",
root.as_ref().user_display()
);
match root.as_ref().try_exists() {
Expand Down
166 changes: 74 additions & 92 deletions crates/uv/src/commands/project/environment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,69 @@ use crate::commands::project::{
use crate::printer::Printer;
use crate::settings::{NetworkSettings, ResolverInstallerSettings};

/// An ephemeral [`PythonEnvironment`] for running an individual command.
#[derive(Debug)]
pub(crate) struct EphemeralEnvironment(PythonEnvironment);

impl From<PythonEnvironment> for EphemeralEnvironment {
fn from(environment: PythonEnvironment) -> Self {
Self(environment)
}
}

impl From<EphemeralEnvironment> for PythonEnvironment {
fn from(environment: EphemeralEnvironment) -> Self {
environment.0
}
}

impl EphemeralEnvironment {
/// Set the ephemeral overlay for a Python environment.
#[allow(clippy::result_large_err)]
pub(crate) fn set_overlay(&self, contents: impl AsRef<[u8]>) -> Result<(), ProjectError> {
let site_packages = self
.0
.site_packages()
.next()
.ok_or(ProjectError::NoSitePackages)?;
let overlay_path = site_packages.join("_uv_ephemeral_overlay.pth");
fs_err::write(overlay_path, contents)?;
Ok(())
}

/// Enable system site packages for a Python environment.
#[allow(clippy::result_large_err)]
pub(crate) fn set_system_site_packages(&self) -> Result<(), ProjectError> {
self.0
.set_pyvenv_cfg("include-system-site-packages", "true")?;
Ok(())
}

/// Set the `extends-environment` key in the `pyvenv.cfg` file to the given path.
///
/// Ephemeral environments created by `uv run --with` extend a parent (virtual or system)
/// environment by adding a `.pth` file to the ephemeral environment's `site-packages`
/// directory. The `pth` file contains Python code to dynamically add the parent
/// environment's `site-packages` directory to Python's import search paths in addition to
/// the ephemeral environment's `site-packages` directory. This works well at runtime, but
/// is too dynamic for static analysis tools like ty to understand. As such, we
/// additionally write the `sys.prefix` of the parent environment to to the
/// `extends-environment` key of the ephemeral environment's `pyvenv.cfg` file, making it
/// easier for these tools to statically and reliably understand the relationship between
/// the two environments.
#[allow(clippy::result_large_err)]
pub(crate) fn set_parent_environment(
&self,
parent_environment_sys_prefix: &Path,
) -> Result<(), ProjectError> {
self.0.set_pyvenv_cfg(
"extends-environment",
&parent_environment_sys_prefix.escape_for_python(),
)?;
Ok(())
}
}

/// A [`PythonEnvironment`] stored in the cache.
#[derive(Debug)]
pub(crate) struct CachedEnvironment(PythonEnvironment);
Expand Down Expand Up @@ -44,15 +107,13 @@ impl CachedEnvironment {
printer: Printer,
preview: PreviewMode,
) -> Result<Self, ProjectError> {
// Resolve the "base" interpreter, which resolves to an underlying parent interpreter if the
// given interpreter is a virtual environment.
let base_interpreter = Self::base_interpreter(interpreter, cache)?;
let interpreter = Self::base_interpreter(interpreter, cache)?;

// Resolve the requirements with the interpreter.
let resolution = Resolution::from(
resolve_environment(
spec,
&base_interpreter,
&interpreter,
build_constraints.clone(),
&settings.resolver,
network_settings,
Expand Down Expand Up @@ -80,29 +141,20 @@ impl CachedEnvironment {
// Use the canonicalized base interpreter path since that's the interpreter we performed the
// resolution with and the interpreter the environment will be created with.
//
// We also include the canonicalized `sys.prefix` of the non-base interpreter, that is, the
// virtual environment's path. Originally, we shared cached environments independent of the
// environment they'd be layered on top of. However, this causes collisions as the overlay
// `.pth` file can be overridden by another instance of uv. Including this element in the key
// avoids this problem at the cost of creating separate cached environments for identical
// `--with` invocations across projects. We use `sys.prefix` rather than `sys.executable` so
// we can canonicalize it without invalidating the purpose of the element — it'd probably be
// safe to just use the absolute `sys.executable` as well.
//
// TODO(zanieb): Since we're not sharing these environmments across projects, we should move
// [`CachedEvnvironment::set_overlay`] etc. here since the values there should be constant
// now.
// We cache environments independent of the environment they'd be layered on top of. The
// assumption is such that the environment will _not_ be modified by the user or uv;
// otherwise, we risk cache poisoning. For example, if we were to write a `.pth` file to
// the cached environment, it would be shared across all projects that use the same
// interpreter and the same cached dependencies.
//
// TODO(zanieb): We should include the version of the base interpreter in the hash, so if
// the interpreter at the canonicalized path changes versions we construct a new
// environment.
let environment_hash = cache_digest(&(
&canonicalize_executable(base_interpreter.sys_executable())?,
&interpreter.sys_prefix().canonicalize()?,
));
let interpreter_hash =
cache_digest(&canonicalize_executable(interpreter.sys_executable())?);

// Search in the content-addressed cache.
let cache_entry = cache.entry(CacheBucket::Environments, environment_hash, resolution_hash);
let cache_entry = cache.entry(CacheBucket::Environments, interpreter_hash, resolution_hash);

if cache.refresh().is_none() {
if let Ok(root) = cache.resolve_link(cache_entry.path()) {
Expand All @@ -116,7 +168,7 @@ impl CachedEnvironment {
let temp_dir = cache.venv_dir()?;
let venv = uv_virtualenv::create_venv(
temp_dir.path(),
base_interpreter,
interpreter,
uv_virtualenv::Prompt::None,
false,
false,
Expand Down Expand Up @@ -150,76 +202,6 @@ impl CachedEnvironment {
Ok(Self(PythonEnvironment::from_root(root, cache)?))
}

/// Set the ephemeral overlay for a Python environment.
#[allow(clippy::result_large_err)]
pub(crate) fn set_overlay(&self, contents: impl AsRef<[u8]>) -> Result<(), ProjectError> {
let site_packages = self
.0
.site_packages()
.next()
.ok_or(ProjectError::NoSitePackages)?;
let overlay_path = site_packages.join("_uv_ephemeral_overlay.pth");
fs_err::write(overlay_path, contents)?;
Ok(())
}

/// Clear the ephemeral overlay for a Python environment, if it exists.
#[allow(clippy::result_large_err)]
pub(crate) fn clear_overlay(&self) -> Result<(), ProjectError> {
let site_packages = self
.0
.site_packages()
.next()
.ok_or(ProjectError::NoSitePackages)?;
let overlay_path = site_packages.join("_uv_ephemeral_overlay.pth");
match fs_err::remove_file(overlay_path) {
Ok(()) => (),
Err(err) if err.kind() == std::io::ErrorKind::NotFound => (),
Err(err) => return Err(ProjectError::OverlayRemoval(err)),
}
Ok(())
}

/// Enable system site packages for a Python environment.
#[allow(clippy::result_large_err)]
pub(crate) fn set_system_site_packages(&self) -> Result<(), ProjectError> {
self.0
.set_pyvenv_cfg("include-system-site-packages", "true")?;
Ok(())
}

/// Disable system site packages for a Python environment.
#[allow(clippy::result_large_err)]
pub(crate) fn clear_system_site_packages(&self) -> Result<(), ProjectError> {
self.0
.set_pyvenv_cfg("include-system-site-packages", "false")?;
Ok(())
}

/// Set the `extends-environment` key in the `pyvenv.cfg` file to the given path.
///
/// Ephemeral environments created by `uv run --with` extend a parent (virtual or system)
/// environment by adding a `.pth` file to the ephemeral environment's `site-packages`
/// directory. The `pth` file contains Python code to dynamically add the parent
/// environment's `site-packages` directory to Python's import search paths in addition to
/// the ephemeral environment's `site-packages` directory. This works well at runtime, but
/// is too dynamic for static analysis tools like ty to understand. As such, we
/// additionally write the `sys.prefix` of the parent environment to the
/// `extends-environment` key of the ephemeral environment's `pyvenv.cfg` file, making it
/// easier for these tools to statically and reliably understand the relationship between
/// the two environments.
#[allow(clippy::result_large_err)]
pub(crate) fn set_parent_environment(
&self,
parent_environment_sys_prefix: &Path,
) -> Result<(), ProjectError> {
self.0.set_pyvenv_cfg(
"extends-environment",
&parent_environment_sys_prefix.escape_for_python(),
)?;
Ok(())
}

/// Return the [`Interpreter`] to use for the cached environment, based on a given
/// [`Interpreter`].
///
Expand Down
3 changes: 0 additions & 3 deletions crates/uv/src/commands/project/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -200,9 +200,6 @@ pub(crate) enum ProjectError {
#[error("Failed to parse PEP 723 script metadata")]
Pep723ScriptTomlParse(#[source] toml::de::Error),

#[error("Failed to remove ephemeral overlay")]
OverlayRemoval(#[source] std::io::Error),

#[error("Failed to find `site-packages` directory for environment")]
NoSitePackages,

Expand Down
120 changes: 82 additions & 38 deletions crates/uv/src/commands/project/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ use crate::commands::pip::loggers::{
DefaultInstallLogger, DefaultResolveLogger, SummaryInstallLogger, SummaryResolveLogger,
};
use crate::commands::pip::operations::Modifications;
use crate::commands::project::environment::CachedEnvironment;
use crate::commands::project::environment::{CachedEnvironment, EphemeralEnvironment};
use crate::commands::project::install_target::InstallTarget;
use crate::commands::project::lock::LockMode;
use crate::commands::project::lock_target::LockTarget;
Expand Down Expand Up @@ -944,15 +944,15 @@ hint: If you are running a script with `{}` in the shebang, you may need to incl

// If necessary, create an environment for the ephemeral requirements or command.
let base_site_packages = SitePackages::from_interpreter(&base_interpreter)?;
let ephemeral_env = match spec {
let requirements_env = match spec {
None => None,
Some(spec)
if can_skip_ephemeral(&spec, &base_interpreter, &base_site_packages, &settings) =>
{
None
}
Some(spec) => {
debug!("Syncing ephemeral requirements");
debug!("Syncing `--with` requirements to cached environment");

// Read the build constraints from the lock file.
let build_constraints = base_lock
Expand Down Expand Up @@ -1013,54 +1013,92 @@ hint: If you are running a script with `{}` in the shebang, you may need to incl
Err(err) => return Err(err.into()),
};

Some(environment)
Some(PythonEnvironment::from(environment))
}
};

// If we're running in an ephemeral environment, add a path file to enable loading of
// the base environment's site packages. Setting `PYTHONPATH` is insufficient, as it doesn't
// resolve `.pth` files in the base environment.
// If we're layering requirements atop the project environment, run the command in an ephemeral,
// isolated environment. Otherwise, modifications to the "active virtual environment" would
// poison the cache.
let ephemeral_dir = requirements_env
.as_ref()
.map(|_| cache.venv_dir())
.transpose()?;

let ephemeral_env = ephemeral_dir
.as_ref()
.map(|dir| {
debug!(
"Creating ephemeral environment at: `{}`",
dir.path().simplified_display()
);

uv_virtualenv::create_venv(
dir.path(),
base_interpreter.clone(),
uv_virtualenv::Prompt::None,
false,
false,
false,
false,
false,
preview,
)
})
.transpose()?
.map(EphemeralEnvironment::from);

// If we're running in an ephemeral environment, add a path file to enable loading from the
// `--with` requirements environment and the project environment site packages.
//
// `sitecustomize.py` would be an alternative, but it can be shadowed by an existing such
// module in the python installation.
// Setting `PYTHONPATH` is insufficient, as it doesn't resolve `.pth` files in the base
// environment. Adding `sitecustomize.py` would be an alternative, but it can be shadowed by an
// existing such module in the python installation.
if let Some(ephemeral_env) = ephemeral_env.as_ref() {
let site_packages = base_interpreter
.site_packages()
.next()
.ok_or_else(|| ProjectError::NoSitePackages)?;
ephemeral_env.set_overlay(format!(
"import site; site.addsitedir(\"{}\")",
site_packages.escape_for_python()
))?;

// 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`).
//
// Note that we do this even if the parent environment is not a virtual environment.
// For ephemeral environments created by `uv run --with`, the parent environment's
// `site-packages` directory is added to `sys.path` even if the parent environment is not
// a virtual environment and even if `--system-site-packages` was not explicitly selected.
ephemeral_env.set_parent_environment(base_interpreter.sys_prefix())?;

// If `--system-site-packages` is enabled, add the system site packages to the ephemeral
// environment.
if base_interpreter.is_virtualenv()
&& PyVenvConfiguration::parse(base_interpreter.sys_prefix().join("pyvenv.cfg"))
.is_ok_and(|cfg| cfg.include_system_site_packages())
{
ephemeral_env.set_system_site_packages()?;
} else {
ephemeral_env.clear_system_site_packages()?;
if let Some(requirements_env) = requirements_env.as_ref() {
let requirements_site_packages =
requirements_env.site_packages().next().ok_or_else(|| {
anyhow!("Requirements environment has no site packages directory")
})?;
let base_site_packages = base_interpreter
.site_packages()
.next()
.ok_or_else(|| anyhow!("Base environment has no site packages directory"))?;

ephemeral_env.set_overlay(format!(
"import site; site.addsitedir(\"{}\"); site.addsitedir(\"{}\");",
base_site_packages.escape_for_python(),
requirements_site_packages.escape_for_python(),
))?;

// 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`).
//
// Note that we do this even if the parent environment is not a virtual environment.
// For ephemeral environments created by `uv run --with`, the parent environment's
// `site-packages` directory is added to `sys.path` even if the parent environment is not
// a virtual environment and even if `--system-site-packages` was not explicitly selected.
ephemeral_env.set_parent_environment(base_interpreter.sys_prefix())?;

// If `--system-site-packages` is enabled, add the system site packages to the ephemeral
// environment.
if base_interpreter.is_virtualenv()
&& PyVenvConfiguration::parse(base_interpreter.sys_prefix().join("pyvenv.cfg"))
.is_ok_and(|cfg| cfg.include_system_site_packages())
{
ephemeral_env.set_system_site_packages()?;
}
}
}

// Cast from `CachedEnvironment` to `PythonEnvironment`.
// Cast to `PythonEnvironment`.
let ephemeral_env = ephemeral_env.map(PythonEnvironment::from);

// Determine the Python interpreter to use for the command, if necessary.
let interpreter = ephemeral_env
.as_ref()
.or(requirements_env.as_ref())
.map_or_else(|| &base_interpreter, |env| env.interpreter());

// Check if any run command is given.
Expand Down Expand Up @@ -1143,6 +1181,12 @@ hint: If you are running a script with `{}` in the shebang, you may need to incl
.as_ref()
.map(PythonEnvironment::scripts)
.into_iter()
.chain(
requirements_env
.as_ref()
.map(PythonEnvironment::scripts)
.into_iter(),
)
.chain(std::iter::once(base_interpreter.scripts()))
.chain(
// On Windows, non-virtual Python distributions put `python.exe` in the top-level
Expand Down
Loading
Loading