Skip to content

Commit ea4ca87

Browse files
committed
Use an extra layer of ephemeral deps
1 parent e6270b7 commit ea4ca87

File tree

7 files changed

+161
-149
lines changed

7 files changed

+161
-149
lines changed

crates/uv-python/src/environment.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ impl PythonEnvironment {
174174
/// N.B. This function also works for system Python environments and users depend on this.
175175
pub fn from_root(root: impl AsRef<Path>, cache: &Cache) -> Result<Self, Error> {
176176
debug!(
177-
"Checking for Python environment at `{}`",
177+
"Checking for Python environment at: `{}`",
178178
root.as_ref().user_display()
179179
);
180180
match root.as_ref().try_exists() {

crates/uv/src/commands/project/environment.rs

Lines changed: 74 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,69 @@ use crate::commands::project::{
1717
use crate::printer::Printer;
1818
use crate::settings::{NetworkSettings, ResolverInstallerSettings};
1919

20+
/// An ephemeral [`PythonEnvironment`] for running an individual command.
21+
#[derive(Debug)]
22+
pub(crate) struct EphemeralEnvironment(PythonEnvironment);
23+
24+
impl From<PythonEnvironment> for EphemeralEnvironment {
25+
fn from(environment: PythonEnvironment) -> Self {
26+
Self(environment)
27+
}
28+
}
29+
30+
impl From<EphemeralEnvironment> for PythonEnvironment {
31+
fn from(environment: EphemeralEnvironment) -> Self {
32+
environment.0
33+
}
34+
}
35+
36+
impl EphemeralEnvironment {
37+
/// Set the ephemeral overlay for a Python environment.
38+
#[allow(clippy::result_large_err)]
39+
pub(crate) fn set_overlay(&self, contents: impl AsRef<[u8]>) -> Result<(), ProjectError> {
40+
let site_packages = self
41+
.0
42+
.site_packages()
43+
.next()
44+
.ok_or(ProjectError::NoSitePackages)?;
45+
let overlay_path = site_packages.join("_uv_ephemeral_overlay.pth");
46+
fs_err::write(overlay_path, contents)?;
47+
Ok(())
48+
}
49+
50+
/// Enable system site packages for a Python environment.
51+
#[allow(clippy::result_large_err)]
52+
pub(crate) fn set_system_site_packages(&self) -> Result<(), ProjectError> {
53+
self.0
54+
.set_pyvenv_cfg("include-system-site-packages", "true")?;
55+
Ok(())
56+
}
57+
58+
/// Set the `extends-environment` key in the `pyvenv.cfg` file to the given path.
59+
///
60+
/// Ephemeral environments created by `uv run --with` extend a parent (virtual or system)
61+
/// environment by adding a `.pth` file to the ephemeral environment's `site-packages`
62+
/// directory. The `pth` file contains Python code to dynamically add the parent
63+
/// environment's `site-packages` directory to Python's import search paths in addition to
64+
/// the ephemeral environment's `site-packages` directory. This works well at runtime, but
65+
/// is too dynamic for static analysis tools like ty to understand. As such, we
66+
/// additionally write the `sys.prefix` of the parent environment to to the
67+
/// `extends-environment` key of the ephemeral environment's `pyvenv.cfg` file, making it
68+
/// easier for these tools to statically and reliably understand the relationship between
69+
/// the two environments.
70+
#[allow(clippy::result_large_err)]
71+
pub(crate) fn set_parent_environment(
72+
&self,
73+
parent_environment_sys_prefix: &Path,
74+
) -> Result<(), ProjectError> {
75+
self.0.set_pyvenv_cfg(
76+
"extends-environment",
77+
&parent_environment_sys_prefix.escape_for_python(),
78+
)?;
79+
Ok(())
80+
}
81+
}
82+
2083
/// A [`PythonEnvironment`] stored in the cache.
2184
#[derive(Debug)]
2285
pub(crate) struct CachedEnvironment(PythonEnvironment);
@@ -44,15 +107,13 @@ impl CachedEnvironment {
44107
printer: Printer,
45108
preview: PreviewMode,
46109
) -> Result<Self, ProjectError> {
47-
// Resolve the "base" interpreter, which resolves to an underlying parent interpreter if the
48-
// given interpreter is a virtual environment.
49-
let base_interpreter = Self::base_interpreter(interpreter, cache)?;
110+
let interpreter = Self::base_interpreter(interpreter, cache)?;
50111

51112
// Resolve the requirements with the interpreter.
52113
let resolution = Resolution::from(
53114
resolve_environment(
54115
spec,
55-
&base_interpreter,
116+
&interpreter,
56117
build_constraints.clone(),
57118
&settings.resolver,
58119
network_settings,
@@ -80,29 +141,20 @@ impl CachedEnvironment {
80141
// Use the canonicalized base interpreter path since that's the interpreter we performed the
81142
// resolution with and the interpreter the environment will be created with.
82143
//
83-
// We also include the canonicalized `sys.prefix` of the non-base interpreter, that is, the
84-
// virtual environment's path. Originally, we shared cached environments independent of the
85-
// environment they'd be layered on top of. However, this causes collisions as the overlay
86-
// `.pth` file can be overridden by another instance of uv. Including this element in the key
87-
// avoids this problem at the cost of creating separate cached environments for identical
88-
// `--with` invocations across projects. We use `sys.prefix` rather than `sys.executable` so
89-
// we can canonicalize it without invalidating the purpose of the element — it'd probably be
90-
// safe to just use the absolute `sys.executable` as well.
91-
//
92-
// TODO(zanieb): Since we're not sharing these environmments across projects, we should move
93-
// [`CachedEvnvironment::set_overlay`] etc. here since the values there should be constant
94-
// now.
144+
// We cache environments independent of the environment they'd be layered on top of. The
145+
// assumption is such that the environment will _not_ be modified by the user or uv;
146+
// otherwise, we risk cache poisoning. For example, if we were to write a `.pth` file to
147+
// the cached environment, it would be shared across all projects that use the same
148+
// interpreter and the same cached dependencies.
95149
//
96150
// TODO(zanieb): We should include the version of the base interpreter in the hash, so if
97151
// the interpreter at the canonicalized path changes versions we construct a new
98152
// environment.
99-
let environment_hash = cache_digest(&(
100-
&canonicalize_executable(base_interpreter.sys_executable())?,
101-
&interpreter.sys_prefix().canonicalize()?,
102-
));
153+
let interpreter_hash =
154+
cache_digest(&canonicalize_executable(interpreter.sys_executable())?);
103155

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

107159
if cache.refresh().is_none() {
108160
if let Ok(root) = cache.resolve_link(cache_entry.path()) {
@@ -116,7 +168,7 @@ impl CachedEnvironment {
116168
let temp_dir = cache.venv_dir()?;
117169
let venv = uv_virtualenv::create_venv(
118170
temp_dir.path(),
119-
base_interpreter,
171+
interpreter,
120172
uv_virtualenv::Prompt::None,
121173
false,
122174
false,
@@ -150,76 +202,6 @@ impl CachedEnvironment {
150202
Ok(Self(PythonEnvironment::from_root(root, cache)?))
151203
}
152204

153-
/// Set the ephemeral overlay for a Python environment.
154-
#[allow(clippy::result_large_err)]
155-
pub(crate) fn set_overlay(&self, contents: impl AsRef<[u8]>) -> Result<(), ProjectError> {
156-
let site_packages = self
157-
.0
158-
.site_packages()
159-
.next()
160-
.ok_or(ProjectError::NoSitePackages)?;
161-
let overlay_path = site_packages.join("_uv_ephemeral_overlay.pth");
162-
fs_err::write(overlay_path, contents)?;
163-
Ok(())
164-
}
165-
166-
/// Clear the ephemeral overlay for a Python environment, if it exists.
167-
#[allow(clippy::result_large_err)]
168-
pub(crate) fn clear_overlay(&self) -> Result<(), ProjectError> {
169-
let site_packages = self
170-
.0
171-
.site_packages()
172-
.next()
173-
.ok_or(ProjectError::NoSitePackages)?;
174-
let overlay_path = site_packages.join("_uv_ephemeral_overlay.pth");
175-
match fs_err::remove_file(overlay_path) {
176-
Ok(()) => (),
177-
Err(err) if err.kind() == std::io::ErrorKind::NotFound => (),
178-
Err(err) => return Err(ProjectError::OverlayRemoval(err)),
179-
}
180-
Ok(())
181-
}
182-
183-
/// Enable system site packages for a Python environment.
184-
#[allow(clippy::result_large_err)]
185-
pub(crate) fn set_system_site_packages(&self) -> Result<(), ProjectError> {
186-
self.0
187-
.set_pyvenv_cfg("include-system-site-packages", "true")?;
188-
Ok(())
189-
}
190-
191-
/// Disable system site packages for a Python environment.
192-
#[allow(clippy::result_large_err)]
193-
pub(crate) fn clear_system_site_packages(&self) -> Result<(), ProjectError> {
194-
self.0
195-
.set_pyvenv_cfg("include-system-site-packages", "false")?;
196-
Ok(())
197-
}
198-
199-
/// Set the `extends-environment` key in the `pyvenv.cfg` file to the given path.
200-
///
201-
/// Ephemeral environments created by `uv run --with` extend a parent (virtual or system)
202-
/// environment by adding a `.pth` file to the ephemeral environment's `site-packages`
203-
/// directory. The `pth` file contains Python code to dynamically add the parent
204-
/// environment's `site-packages` directory to Python's import search paths in addition to
205-
/// the ephemeral environment's `site-packages` directory. This works well at runtime, but
206-
/// is too dynamic for static analysis tools like ty to understand. As such, we
207-
/// additionally write the `sys.prefix` of the parent environment to the
208-
/// `extends-environment` key of the ephemeral environment's `pyvenv.cfg` file, making it
209-
/// easier for these tools to statically and reliably understand the relationship between
210-
/// the two environments.
211-
#[allow(clippy::result_large_err)]
212-
pub(crate) fn set_parent_environment(
213-
&self,
214-
parent_environment_sys_prefix: &Path,
215-
) -> Result<(), ProjectError> {
216-
self.0.set_pyvenv_cfg(
217-
"extends-environment",
218-
&parent_environment_sys_prefix.escape_for_python(),
219-
)?;
220-
Ok(())
221-
}
222-
223205
/// Return the [`Interpreter`] to use for the cached environment, based on a given
224206
/// [`Interpreter`].
225207
///

crates/uv/src/commands/project/mod.rs

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -200,9 +200,6 @@ pub(crate) enum ProjectError {
200200
#[error("Failed to parse PEP 723 script metadata")]
201201
Pep723ScriptTomlParse(#[source] toml::de::Error),
202202

203-
#[error("Failed to remove ephemeral overlay")]
204-
OverlayRemoval(#[source] std::io::Error),
205-
206203
#[error("Failed to find `site-packages` directory for environment")]
207204
NoSitePackages,
208205

crates/uv/src/commands/project/run.rs

Lines changed: 82 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ use crate::commands::pip::loggers::{
4545
DefaultInstallLogger, DefaultResolveLogger, SummaryInstallLogger, SummaryResolveLogger,
4646
};
4747
use crate::commands::pip::operations::Modifications;
48-
use crate::commands::project::environment::CachedEnvironment;
48+
use crate::commands::project::environment::{CachedEnvironment, EphemeralEnvironment};
4949
use crate::commands::project::install_target::InstallTarget;
5050
use crate::commands::project::lock::LockMode;
5151
use crate::commands::project::lock_target::LockTarget;
@@ -944,15 +944,15 @@ hint: If you are running a script with `{}` in the shebang, you may need to incl
944944

945945
// If necessary, create an environment for the ephemeral requirements or command.
946946
let base_site_packages = SitePackages::from_interpreter(&base_interpreter)?;
947-
let ephemeral_env = match spec {
947+
let requirements_env = match spec {
948948
None => None,
949949
Some(spec)
950950
if can_skip_ephemeral(&spec, &base_interpreter, &base_site_packages, &settings) =>
951951
{
952952
None
953953
}
954954
Some(spec) => {
955-
debug!("Syncing ephemeral requirements");
955+
debug!("Syncing `--with` requirements to cached environment");
956956

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

1016-
Some(environment)
1016+
Some(PythonEnvironment::from(environment))
10171017
}
10181018
};
10191019

1020-
// If we're running in an ephemeral environment, add a path file to enable loading of
1021-
// the base environment's site packages. Setting `PYTHONPATH` is insufficient, as it doesn't
1022-
// resolve `.pth` files in the base environment.
1020+
// If we're layering requirements atop the project environment, run the command in an ephemeral,
1021+
// isolated environment. Otherwise, modifications to the "active virtual environment" would
1022+
// poison the cache.
1023+
let ephemeral_dir = requirements_env
1024+
.as_ref()
1025+
.map(|_| cache.venv_dir())
1026+
.transpose()?;
1027+
1028+
let ephemeral_env = ephemeral_dir
1029+
.as_ref()
1030+
.map(|dir| {
1031+
debug!(
1032+
"Creating ephemeral environment at: `{}`",
1033+
dir.path().simplified_display()
1034+
);
1035+
1036+
uv_virtualenv::create_venv(
1037+
dir.path(),
1038+
base_interpreter.clone(),
1039+
uv_virtualenv::Prompt::None,
1040+
false,
1041+
false,
1042+
false,
1043+
false,
1044+
false,
1045+
preview,
1046+
)
1047+
})
1048+
.transpose()?
1049+
.map(EphemeralEnvironment::from);
1050+
1051+
// If we're running in an ephemeral environment, add a path file to enable loading from the
1052+
// `--with` requirements environment and the project environment site packages.
10231053
//
1024-
// `sitecustomize.py` would be an alternative, but it can be shadowed by an existing such
1025-
// module in the python installation.
1054+
// Setting `PYTHONPATH` is insufficient, as it doesn't resolve `.pth` files in the base
1055+
// environment. Adding `sitecustomize.py` would be an alternative, but it can be shadowed by an
1056+
// existing such module in the python installation.
10261057
if let Some(ephemeral_env) = ephemeral_env.as_ref() {
1027-
let site_packages = base_interpreter
1028-
.site_packages()
1029-
.next()
1030-
.ok_or_else(|| ProjectError::NoSitePackages)?;
1031-
ephemeral_env.set_overlay(format!(
1032-
"import site; site.addsitedir(\"{}\")",
1033-
site_packages.escape_for_python()
1034-
))?;
1035-
1036-
// Write the `sys.prefix` of the parent environment to the `extends-environment` key of the `pyvenv.cfg`
1037-
// file. This helps out static-analysis tools such as ty (see docs on
1038-
// `CachedEnvironment::set_parent_environment`).
1039-
//
1040-
// Note that we do this even if the parent environment is not a virtual environment.
1041-
// For ephemeral environments created by `uv run --with`, the parent environment's
1042-
// `site-packages` directory is added to `sys.path` even if the parent environment is not
1043-
// a virtual environment and even if `--system-site-packages` was not explicitly selected.
1044-
ephemeral_env.set_parent_environment(base_interpreter.sys_prefix())?;
1045-
1046-
// If `--system-site-packages` is enabled, add the system site packages to the ephemeral
1047-
// environment.
1048-
if base_interpreter.is_virtualenv()
1049-
&& PyVenvConfiguration::parse(base_interpreter.sys_prefix().join("pyvenv.cfg"))
1050-
.is_ok_and(|cfg| cfg.include_system_site_packages())
1051-
{
1052-
ephemeral_env.set_system_site_packages()?;
1053-
} else {
1054-
ephemeral_env.clear_system_site_packages()?;
1058+
if let Some(requirements_env) = requirements_env.as_ref() {
1059+
let requirements_site_packages =
1060+
requirements_env.site_packages().next().ok_or_else(|| {
1061+
anyhow!("Requirements environment has no site packages directory")
1062+
})?;
1063+
let base_site_packages = base_interpreter
1064+
.site_packages()
1065+
.next()
1066+
.ok_or_else(|| anyhow!("Base environment has no site packages directory"))?;
1067+
1068+
ephemeral_env.set_overlay(format!(
1069+
"import site; site.addsitedir(\"{}\"); site.addsitedir(\"{}\");",
1070+
base_site_packages.escape_for_python(),
1071+
requirements_site_packages.escape_for_python(),
1072+
))?;
1073+
1074+
// Write the `sys.prefix` of the parent environment to the `extends-environment` key of the `pyvenv.cfg`
1075+
// file. This helps out static-analysis tools such as ty (see docs on
1076+
// `CachedEnvironment::set_parent_environment`).
1077+
//
1078+
// Note that we do this even if the parent environment is not a virtual environment.
1079+
// For ephemeral environments created by `uv run --with`, the parent environment's
1080+
// `site-packages` directory is added to `sys.path` even if the parent environment is not
1081+
// a virtual environment and even if `--system-site-packages` was not explicitly selected.
1082+
ephemeral_env.set_parent_environment(base_interpreter.sys_prefix())?;
1083+
1084+
// If `--system-site-packages` is enabled, add the system site packages to the ephemeral
1085+
// environment.
1086+
if base_interpreter.is_virtualenv()
1087+
&& PyVenvConfiguration::parse(base_interpreter.sys_prefix().join("pyvenv.cfg"))
1088+
.is_ok_and(|cfg| cfg.include_system_site_packages())
1089+
{
1090+
ephemeral_env.set_system_site_packages()?;
1091+
}
10551092
}
10561093
}
10571094

1058-
// Cast from `CachedEnvironment` to `PythonEnvironment`.
1095+
// Cast to `PythonEnvironment`.
10591096
let ephemeral_env = ephemeral_env.map(PythonEnvironment::from);
10601097

10611098
// Determine the Python interpreter to use for the command, if necessary.
10621099
let interpreter = ephemeral_env
10631100
.as_ref()
1101+
.or(requirements_env.as_ref())
10641102
.map_or_else(|| &base_interpreter, |env| env.interpreter());
10651103

10661104
// Check if any run command is given.
@@ -1143,6 +1181,12 @@ hint: If you are running a script with `{}` in the shebang, you may need to incl
11431181
.as_ref()
11441182
.map(PythonEnvironment::scripts)
11451183
.into_iter()
1184+
.chain(
1185+
requirements_env
1186+
.as_ref()
1187+
.map(PythonEnvironment::scripts)
1188+
.into_iter(),
1189+
)
11461190
.chain(std::iter::once(base_interpreter.scripts()))
11471191
.chain(
11481192
// On Windows, non-virtual Python distributions put `python.exe` in the top-level

0 commit comments

Comments
 (0)