Use stable environments for remote and stdin scripts

astral-sh · Feb 11, 2025 · 5920b63 · 5920b63
1 parent ad3294b
commit 5920b63
Show file tree

Hide file tree

Showing 8 changed files with 157 additions and 197 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/uv-python/src/environment.rs b/crates/uv-python/src/environment.rs
@@ -1,10 +1,12 @@
-use owo_colors::OwoColorize;
 use std::borrow::Cow;
 use std::env;
 use std::fmt;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
+
+use owo_colors::OwoColorize;
 use tracing::debug;
+
 use uv_cache::Cache;
 use uv_cache_key::cache_digest;
 use uv_fs::{LockedFile, Simplified};

diff --git a/crates/uv-scripts/Cargo.toml b/crates/uv-scripts/Cargo.toml
@@ -23,3 +23,4 @@ memchr = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 thiserror = { workspace = true }
 toml = { workspace = true }
+url = { workspace = true }
diff --git a/crates/uv-scripts/src/lib.rs b/crates/uv-scripts/src/lib.rs
@@ -7,6 +7,7 @@ use std::sync::LazyLock;
 use memchr::memmem::Finder;
 use serde::Deserialize;
 use thiserror::Error;
+use url::Url;
 
 use uv_pep440::VersionSpecifiers;
 use uv_pep508::PackageName;
@@ -24,7 +25,7 @@ pub enum Pep723Item {
     /// A PEP 723 script provided via `stdin`.
     Stdin(Pep723Metadata),
     /// A PEP 723 script provided via a remote URL.
-    Remote(Pep723Metadata),
+    Remote(Pep723Metadata, Url),
 }
 
 impl Pep723Item {
@@ -33,7 +34,7 @@ impl Pep723Item {
         match self {
             Self::Script(script) => &script.metadata,
             Self::Stdin(metadata) => metadata,
-            Self::Remote(metadata) => metadata,
+            Self::Remote(metadata, ..) => metadata,
         }
     }
 
@@ -42,16 +43,16 @@ impl Pep723Item {
         match self {
             Self::Script(script) => script.metadata,
             Self::Stdin(metadata) => metadata,
-            Self::Remote(metadata) => metadata,
+            Self::Remote(metadata, ..) => metadata,
         }
     }
 
     /// Return the path of the PEP 723 item, if any.
     pub fn path(&self) -> Option<&Path> {
         match self {
             Self::Script(script) => Some(&script.path),
-            Self::Stdin(_) => None,
-            Self::Remote(_) => None,
+            Self::Stdin(..) => None,
+            Self::Remote(..) => None,
         }
     }
 
@@ -72,7 +73,7 @@ pub enum Pep723ItemRef<'item> {
     /// A PEP 723 script provided via `stdin`.
     Stdin(&'item Pep723Metadata),
     /// A PEP 723 script provided via a remote URL.
-    Remote(&'item Pep723Metadata),
+    Remote(&'item Pep723Metadata, &'item Url),
 }
 
 impl Pep723ItemRef<'_> {
@@ -81,16 +82,16 @@ impl Pep723ItemRef<'_> {
         match self {
             Self::Script(script) => &script.metadata,
             Self::Stdin(metadata) => metadata,
-            Self::Remote(metadata) => metadata,
+            Self::Remote(metadata, ..) => metadata,
         }
     }
 
     /// Return the path of the PEP 723 item, if any.
     pub fn path(&self) -> Option<&Path> {
         match self {
             Self::Script(script) => Some(&script.path),
-            Self::Stdin(_) => None,
-            Self::Remote(_) => None,
+            Self::Stdin(..) => None,
+            Self::Remote(..) => None,
         }
     }
 }
@@ -100,7 +101,7 @@ impl<'item> From<&'item Pep723Item> for Pep723ItemRef<'item> {
         match item {
             Pep723Item::Script(script) => Self::Script(script),
             Pep723Item::Stdin(metadata) => Self::Stdin(metadata),
-            Pep723Item::Remote(metadata) => Self::Remote(metadata),
+            Pep723Item::Remote(metadata, url) => Self::Remote(metadata, url),
         }
     }
 }

diff --git a/crates/uv/src/commands/project/environment.rs b/crates/uv/src/commands/project/environment.rs
@@ -157,11 +157,6 @@ impl CachedEnvironment {
         Ok(())
     }
 
-    /// Convert the [`CachedEnvironment`] into an [`Interpreter`].
-    pub(crate) fn into_interpreter(self) -> Interpreter {
-        self.0.into_interpreter()
-    }
-
     /// Return the [`Interpreter`] to use for the cached environment, based on a given
     /// [`Interpreter`].
     ///

diff --git a/crates/uv/src/commands/project/mod.rs b/crates/uv/src/commands/project/mod.rs
@@ -37,7 +37,7 @@ use uv_resolver::{
     FlatIndex, Lock, OptionsBuilder, PythonRequirement, RequiresPython, ResolverEnvironment,
     ResolverOutput,
 };
-use uv_scripts::{Pep723ItemRef, Pep723Script};
+use uv_scripts::Pep723ItemRef;
 use uv_settings::PythonInstallMirrors;
 use uv_types::{BuildIsolation, EmptyInstalledPackages, HashStrategy};
 use uv_warnings::{warn_user, warn_user_once};
@@ -525,14 +525,21 @@ pub(crate) enum ScriptInterpreter {
 
 impl ScriptInterpreter {
     /// Return the expected virtual environment path for the [`Pep723Script`].
-    pub(crate) fn root(script: &Pep723Script, cache: &Cache) -> PathBuf {
-        let digest = cache_digest(&script.path);
-        let entry = if let Some(file_name) = script.path.file_stem().and_then(|name| name.to_str())
-        {
-            // Replace any non-ASCII characters with underscores.
-            format!("{file_name}-{digest}",)
-        } else {
-            digest
+    pub(crate) fn root(script: Pep723ItemRef<'_>, cache: &Cache) -> PathBuf {
+        let entry = match script {
+            // For local scripts, use a hash of the path to the script.
+            Pep723ItemRef::Script(script) => {
+                let digest = cache_digest(&script.path);
+                if let Some(file_name) = script.path.file_stem().and_then(|name| name.to_str()) {
+                    format!("{file_name}-{digest}")
+                } else {
+                    digest
+                }
+            }
+            // For remote scripts, use a hash of the URL.
+            Pep723ItemRef::Remote(.., url) => cache_digest(url),
+            // Otherwise, use a hash of the metadata.
+            Pep723ItemRef::Stdin(metadata) => cache_digest(&metadata.raw),
         };
         cache
             .shard(CacheBucket::Environments, entry)
@@ -562,42 +569,39 @@ impl ScriptInterpreter {
             requires_python,
         } = ScriptPython::from_request(python_request, workspace, script, no_config).await?;
 
-        // If this is a local script, use a stable virtual environment.
-        if let Pep723ItemRef::Script(script) = script {
-            let root = Self::root(script, cache);
-            match PythonEnvironment::from_root(&root, cache) {
-                Ok(venv) => {
-                    if python_request.as_ref().map_or(true, |request| {
-                        if request.satisfied(venv.interpreter(), cache) {
-                            debug!(
-                                "The script environment's Python version satisfies `{}`",
-                                request.to_canonical_string()
-                            );
-                            true
-                        } else {
-                            debug!(
-                                "The script environment's Python version does not satisfy `{}`",
-                                request.to_canonical_string()
-                            );
-                            false
-                        }
-                    }) {
-                        if let Some((requires_python, ..)) = requires_python.as_ref() {
-                            if requires_python.contains(venv.interpreter().python_version()) {
-                                return Ok(Self::Environment(venv));
-                            }
-                            debug!(
-                                "The script environment's Python version does not meet the script's Python requirement: `{requires_python}`"
-                            );
-                        } else {
+        let root = Self::root(script, cache);
+        match PythonEnvironment::from_root(&root, cache) {
+            Ok(venv) => {
+                if python_request.as_ref().map_or(true, |request| {
+                    if request.satisfied(venv.interpreter(), cache) {
+                        debug!(
+                            "The script environment's Python version satisfies `{}`",
+                            request.to_canonical_string()
+                        );
+                        true
+                    } else {
+                        debug!(
+                            "The script environment's Python version does not satisfy `{}`",
+                            request.to_canonical_string()
+                        );
+                        false
+                    }
+                }) {
+                    if let Some((requires_python, ..)) = requires_python.as_ref() {
+                        if requires_python.contains(venv.interpreter().python_version()) {
                             return Ok(Self::Environment(venv));
                         }
+                        debug!(
+                            "The script environment's Python version does not meet the script's Python requirement: `{requires_python}`"
+                        );
+                    } else {
+                        return Ok(Self::Environment(venv));
                     }
                 }
-                Err(uv_python::Error::MissingEnvironment(_)) => {}
-                Err(err) => warn!("Ignoring existing script environment: {err}"),
-            };
-        }
+            }
+            Err(uv_python::Error::MissingEnvironment(_)) => {}
+            Err(err) => warn!("Ignoring existing script environment: {err}"),
+        };
 
         let client_builder = BaseClientBuilder::new()
             .connectivity(connectivity)
@@ -644,12 +648,30 @@ impl ScriptInterpreter {
     }
 
     /// Grab a file lock for the script to prevent concurrent writes across processes.
-    pub(crate) async fn lock(script: &Pep723Script) -> Result<LockedFile, std::io::Error> {
-        LockedFile::acquire(
-            std::env::temp_dir().join(format!("uv-{}.lock", cache_digest(&script.path))),
-            script.path.simplified_display(),
-        )
-        .await
+    pub(crate) async fn lock(script: Pep723ItemRef<'_>) -> Result<LockedFile, std::io::Error> {
+        match script {
+            Pep723ItemRef::Script(script) => {
+                LockedFile::acquire(
+                    std::env::temp_dir().join(format!("uv-{}.lock", cache_digest(&script.path))),
+                    script.path.simplified_display(),
+                )
+                .await
+            }
+            Pep723ItemRef::Remote(.., url) => {
+                LockedFile::acquire(
+                    std::env::temp_dir().join(format!("uv-{}.lock", cache_digest(url))),
+                    url.to_string(),
+                )
+                .await
+            }
+            Pep723ItemRef::Stdin(metadata) => {
+                LockedFile::acquire(
+                    std::env::temp_dir().join(format!("uv-{}.lock", cache_digest(&metadata.raw))),
+                    "stdin".to_string(),
+                )
+                .await
+            }
+        }
     }
 }
 
@@ -1175,7 +1197,7 @@ struct ScriptEnvironment(PythonEnvironment);
 impl ScriptEnvironment {
     /// Initialize a virtual environment for a PEP 723 script.
     pub(crate) async fn get_or_init(
-        script: &Pep723Script,
+        script: Pep723ItemRef<'_>,
         python_request: Option<PythonRequest>,
         python_preference: PythonPreference,
         python_downloads: PythonDownloads,
@@ -1191,7 +1213,7 @@ impl ScriptEnvironment {
         let _lock = ScriptInterpreter::lock(script).await?;
 
         match ScriptInterpreter::discover(
-            Pep723ItemRef::Script(script),
+            script,
             python_request,
             python_preference,
             python_downloads,
@@ -1234,8 +1256,8 @@ impl ScriptEnvironment {
                 // 1) The name of the script
                 // 2) No prompt
                 let prompt = script
-                    .path
-                    .file_name()
+                    .path()
+                    .and_then(|path| path.file_name())
                     .map(|f| f.to_string_lossy().to_string())
                     .map(uv_virtualenv::Prompt::Static)
                     .unwrap_or(uv_virtualenv::Prompt::None);