Write the path of the parent environment to an extends-environment key in the pyvenv.cfg file of an ephemeral environment (#13598)

Author: AlexWaygood

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

@@ -1,9 +1,12 @@
+use std::path::Path;
+
 use tracing::debug;
 
 use uv_cache::{Cache, CacheBucket};
 use uv_cache_key::{cache_digest, hash_digest};
 use uv_configuration::{Concurrency, Constraints, PreviewMode};
 use uv_distribution_types::{Name, Resolution};
+use uv_fs::PythonExt;
 use uv_python::{Interpreter, PythonEnvironment};
 
 use crate::commands::pip::loggers::{InstallLogger, ResolveLogger};
@@ -168,6 +171,30 @@ impl CachedEnvironment {
         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(())
+    }
+
     /// Return the [`Interpreter`] to use for the cached environment, based on a given
     /// [`Interpreter`].
     ///

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

@@ -992,6 +992,16 @@ hint: If you are running a script with `{}` in the shebang, you may need to incl
             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()

crates/uv/tests/it/common/mod.rs

@@ -243,6 +243,30 @@ impl TestContext {
         self
     }
 
+    /// Filtering for various keys in a `pyvenv.cfg` file that will vary
+    /// depending on the specific machine used:
+    /// - `home = foo/bar/baz/python3.X.X/bin`
+    /// - `uv = X.Y.Z`
+    /// - `extends-environment = <path/to/parent/venv>`
+    #[must_use]
+    pub fn with_pyvenv_cfg_filters(mut self) -> Self {
+        let added_filters = [
+            (r"home = .+".to_string(), "home = [PYTHON_HOME]".to_string()),
+            (
+                r"uv = \d+\.\d+\.\d+".to_string(),
+                "uv = [UV_VERSION]".to_string(),
+            ),
+            (
+                r"extends-environment = .+".to_string(),
+                "extends-environment = [PARENT_VENV]".to_string(),
+            ),
+        ];
+        for filter in added_filters {
+            self.filters.insert(0, filter);
+        }
+        self
+    }
+
     /// Add extra filtering for ` -> <PATH>` symlink display for Python versions in the test
     /// context, e.g., for use in `uv python list`.
     #[must_use]

crates/uv/tests/it/run.rs

@@ -1264,6 +1264,62 @@ fn run_with() -> Result<()> {
     Ok(())
 }
 
+/// Test that an ephemeral environment writes the path of its parent environment to the `extends-environment` key
+/// of its `pyvenv.cfg` file. This feature makes it easier for static-analysis tools like ty to resolve which import
+/// search paths are available in these ephemeral environments.
+#[test]
+fn run_with_pyvenv_cfg_file() -> Result<()> {
+    let context = TestContext::new("3.12").with_pyvenv_cfg_filters();
+
+    let pyproject_toml = context.temp_dir.child("pyproject.toml");
+    pyproject_toml.write_str(indoc! { r#"
+        [project]
+        name = "foo"
+        version = "1.0.0"
+        requires-python = ">=3.8"
+
+        [build-system]
+        requires = ["setuptools>=42"]
+        build-backend = "setuptools.build_meta"
+        "#
+    })?;
+
+    let test_script = context.temp_dir.child("main.py");
+    test_script.write_str(indoc! { r#"
+        import os
+
+        with open(f'{os.getenv("VIRTUAL_ENV")}/pyvenv.cfg') as f:
+            print(f.read())
+       "#
+    })?;
+
+    uv_snapshot!(context.filters(), context.run().arg("--with").arg("iniconfig").arg("main.py"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    home = [PYTHON_HOME]
+    implementation = CPython
+    uv = [UV_VERSION]
+    version_info = 3.12.[X]
+    include-system-site-packages = false
+    relocatable = true
+    extends-environment = [PARENT_VENV]
+
+
+    ----- stderr -----
+    Resolved 1 package in [TIME]
+    Prepared 1 package in [TIME]
+    Installed 1 package in [TIME]
+     + foo==1.0.0 (from file://[TEMP_DIR]/)
+    Resolved 1 package in [TIME]
+    Prepared 1 package in [TIME]
+    Installed 1 package in [TIME]
+     + iniconfig==2.0.0
+    ");
+
+    Ok(())
+}
+
 #[test]
 fn run_with_build_constraints() -> Result<()> {
     let context = TestContext::new("3.8");