From d864e1dbe5e3ea0cefa43f89b7e900f69a77c32b Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 11 Oct 2024 00:49:14 +0200 Subject: [PATCH] Use separate types for stdin vs. file-based PEP 723 scripts (#8113) --- crates/uv-scripts/src/lib.rs | 85 ++++++++++++++++-------- crates/uv/src/commands/project/add.rs | 13 ++-- crates/uv/src/commands/project/remove.rs | 8 ++- crates/uv/src/commands/project/run.rs | 35 +++++----- crates/uv/src/lib.rs | 20 ++++-- 5 files changed, 99 insertions(+), 62 deletions(-) diff --git a/crates/uv-scripts/src/lib.rs b/crates/uv-scripts/src/lib.rs index 7ea275c6c..a463da455 100644 --- a/crates/uv-scripts/src/lib.rs +++ b/crates/uv-scripts/src/lib.rs @@ -16,11 +16,38 @@ use uv_workspace::pyproject::Sources; static FINDER: LazyLock = LazyLock::new(|| Finder::new(b"# /// script")); +/// A PEP 723 item, either read from a script on disk or provided via `stdin`. +#[derive(Debug)] +pub enum Pep723Item { + /// A PEP 723 script read from disk. + Script(Pep723Script), + /// A PEP 723 script provided via `stdin`. + Stdin(Pep723Stdin), +} + +impl Pep723Item { + /// Return the [`Pep723Metadata`] associated with the item. + pub fn metadata(&self) -> &Pep723Metadata { + match self { + Self::Script(script) => &script.metadata, + Self::Stdin(stdin) => &stdin.metadata, + } + } + + /// Consume the item and return the associated [`Pep723Metadata`]. + pub fn into_metadata(self) -> Pep723Metadata { + match self { + Self::Script(script) => script.metadata, + Self::Stdin(stdin) => stdin.metadata, + } + } +} + /// A PEP 723 script, including its [`Pep723Metadata`]. #[derive(Debug)] pub struct Pep723Script { /// The path to the Python script. - pub source: Source, + pub path: PathBuf, /// The parsed [`Pep723Metadata`] table from the script. pub metadata: Pep723Metadata, /// The content of the script before the metadata table. @@ -34,28 +61,18 @@ impl Pep723Script { /// /// See: pub async fn read(file: impl AsRef) -> Result, Pep723Error> { - match fs_err::tokio::read(&file).await { - Ok(contents) => { - Self::parse_contents(&contents, Source::File(file.as_ref().to_path_buf())) - } - Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None), - Err(err) => Err(err.into()), - } - } + let contents = match fs_err::tokio::read(&file).await { + Ok(contents) => contents, + Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(None), + Err(err) => return Err(err.into()), + }; - /// Read the PEP 723 `script` metadata from stdin. - pub fn parse_stdin(contents: &[u8]) -> Result, Pep723Error> { - Self::parse_contents(contents, Source::Stdin) - } - - /// Parse the contents of a Python script and extract the `script` metadata block. - fn parse_contents(contents: &[u8], source: Source) -> Result, Pep723Error> { // Extract the `script` tag. let ScriptTag { prelude, metadata, postlude, - } = match ScriptTag::parse(contents) { + } = match ScriptTag::parse(&contents) { Ok(Some(tag)) => tag, Ok(None) => return Ok(None), Err(err) => return Err(err), @@ -65,7 +82,7 @@ impl Pep723Script { let metadata = Pep723Metadata::from_str(&metadata)?; Ok(Some(Self { - source, + path: file.as_ref().to_path_buf(), metadata, prelude, postlude, @@ -94,7 +111,7 @@ impl Pep723Script { let (shebang, postlude) = extract_shebang(&contents)?; Ok(Self { - source: Source::File(file.as_ref().to_path_buf()), + path: file.as_ref().to_path_buf(), prelude: if shebang.is_empty() { String::new() } else { @@ -159,21 +176,33 @@ impl Pep723Script { self.postlude ); - if let Source::File(path) = &self.source { - fs_err::tokio::write(&path, content).await?; - } + fs_err::tokio::write(&self.path, content).await?; Ok(()) } } -/// The source of a PEP 723 script. +/// A PEP 723 script, provided via `stdin`. #[derive(Debug)] -pub enum Source { - /// The PEP 723 script is sourced from a file. - File(PathBuf), - /// The PEP 723 script is sourced from stdin. - Stdin, +pub struct Pep723Stdin { + metadata: Pep723Metadata, +} + +impl Pep723Stdin { + /// Parse the PEP 723 `script` metadata from `stdin`. + pub fn parse(contents: &[u8]) -> Result, Pep723Error> { + // Extract the `script` tag. + let ScriptTag { metadata, .. } = match ScriptTag::parse(contents) { + Ok(Some(tag)) => tag, + Ok(None) => return Ok(None), + Err(err) => return Err(err), + }; + + // Parse the metadata. + let metadata = Pep723Metadata::from_str(&metadata)?; + + Ok(Some(Self { metadata })) + } } /// PEP 723 metadata as parsed from a `script` comment block. diff --git a/crates/uv/src/commands/project/add.rs b/crates/uv/src/commands/project/add.rs index b5f83bc31..43f57d48d 100644 --- a/crates/uv/src/commands/project/add.rs +++ b/crates/uv/src/commands/project/add.rs @@ -379,10 +379,7 @@ pub(crate) async fn add( (uv_pep508::Requirement::from(requirement), None) } Target::Script(ref script, _) => { - let uv_scripts::Source::File(path) = &script.source else { - unreachable!("script source is not a file"); - }; - let script_path = std::path::absolute(path)?; + let script_path = std::path::absolute(&script.path)?; let script_dir = script_path.parent().expect("script path has no parent"); resolve_requirement( requirement, @@ -511,9 +508,11 @@ pub(crate) async fn add( Target::Project(project, venv) => (project, venv), // If `--script`, exit early. There's no reason to lock and sync. Target::Script(script, _) => { - if let uv_scripts::Source::File(path) = &script.source { - writeln!(printer.stderr(), "Updated `{}`", path.user_display().cyan())?; - } + writeln!( + printer.stderr(), + "Updated `{}`", + script.path.user_display().cyan() + )?; return Ok(ExitStatus::Success); } }; diff --git a/crates/uv/src/commands/project/remove.rs b/crates/uv/src/commands/project/remove.rs index 9a2c69906..fd0155c77 100644 --- a/crates/uv/src/commands/project/remove.rs +++ b/crates/uv/src/commands/project/remove.rs @@ -144,9 +144,11 @@ pub(crate) async fn remove( Target::Project(project) => project, // If `--script`, exit early. There's no reason to lock and sync. Target::Script(script) => { - if let uv_scripts::Source::File(path) = &script.source { - writeln!(printer.stderr(), "Updated `{}`", path.user_display().cyan())?; - } + writeln!( + printer.stderr(), + "Updated `{}`", + script.path.user_display().cyan() + )?; return Ok(ExitStatus::Success); } }; diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs index da90fa2f1..bf36eb8ab 100644 --- a/crates/uv/src/commands/project/run.rs +++ b/crates/uv/src/commands/project/run.rs @@ -29,7 +29,7 @@ use uv_python::{ }; use uv_requirements::{RequirementsSource, RequirementsSpecification}; use uv_resolver::Lock; -use uv_scripts::Pep723Script; +use uv_scripts::Pep723Item; use uv_warnings::warn_user; use uv_workspace::{DiscoveryOptions, InstallTarget, VirtualProject, Workspace, WorkspaceError}; @@ -52,7 +52,7 @@ use crate::settings::ResolverInstallerSettings; #[allow(clippy::fn_params_excessive_bools)] pub(crate) async fn run( project_dir: &Path, - script: Option, + script: Option, command: Option, requirements: Vec, show_resolution: bool, @@ -107,11 +107,11 @@ pub(crate) async fn run( // Determine whether the command to execute is a PEP 723 script. let temp_dir; let script_interpreter = if let Some(script) = script { - if let uv_scripts::Source::File(path) = &script.source { + if let Pep723Item::Script(script) = &script { writeln!( printer.stderr(), "Reading inline script metadata from: {}", - path.user_display().cyan() + script.path.user_display().cyan() )?; } else { writeln!( @@ -134,7 +134,7 @@ pub(crate) async fn run( } else { // (3) `Requires-Python` in the script let request = script - .metadata + .metadata() .requires_python .as_ref() .map(|requires_python| { @@ -163,7 +163,7 @@ pub(crate) async fn run( .await? .into_interpreter(); - if let Some(requires_python) = script.metadata.requires_python.as_ref() { + if let Some(requires_python) = script.metadata().requires_python.as_ref() { if !requires_python.contains(interpreter.python_version()) { let err = match source { PythonRequestSource::UserRequest => { @@ -190,13 +190,22 @@ pub(crate) async fn run( } } + // Determine the working directory for the script. + let script_dir = match &script { + Pep723Item::Script(script) => std::path::absolute(&script.path)? + .parent() + .expect("script path has no parent") + .to_owned(), + Pep723Item::Stdin(..) => std::env::current_dir()?, + }; + let script = script.into_metadata(); + // Install the script requirements, if necessary. Otherwise, use an isolated environment. - if let Some(dependencies) = script.metadata.dependencies { + if let Some(dependencies) = script.dependencies { // // Collect any `tool.uv.sources` from the script. let empty = BTreeMap::default(); let script_sources = match settings.sources { SourceStrategy::Enabled => script - .metadata .tool .as_ref() .and_then(|tool| tool.uv.as_ref()) @@ -204,16 +213,6 @@ pub(crate) async fn run( .unwrap_or(&empty), SourceStrategy::Disabled => &empty, }; - let script_dir = match &script.source { - uv_scripts::Source::File(path) => { - let script_path = std::path::absolute(path)?; - script_path - .parent() - .expect("script path has no parent") - .to_owned() - } - uv_scripts::Source::Stdin => std::env::current_dir()?, - }; let requirements = dependencies .into_iter() diff --git a/crates/uv/src/lib.rs b/crates/uv/src/lib.rs index 056fa260e..200fecfdb 100644 --- a/crates/uv/src/lib.rs +++ b/crates/uv/src/lib.rs @@ -28,7 +28,7 @@ use uv_cli::{SelfCommand, SelfNamespace, SelfUpdateArgs}; use uv_client::BaseClientBuilder; use uv_fs::CWD; use uv_requirements::RequirementsSource; -use uv_scripts::Pep723Script; +use uv_scripts::{Pep723Item, Pep723Script, Pep723Stdin}; use uv_settings::{Combine, FilesystemOptions, Options}; use uv_warnings::{warn_user, warn_user_once}; use uv_workspace::{DiscoveryOptions, Workspace}; @@ -217,8 +217,10 @@ async fn run(mut cli: Cli) -> Result { match run_command.as_ref() { Some( RunCommand::PythonScript(script, _) | RunCommand::PythonGuiScript(script, _), - ) => Pep723Script::read(&script).await?, - Some(RunCommand::PythonStdin(contents)) => Pep723Script::parse_stdin(contents)?, + ) => Pep723Script::read(&script).await?.map(Pep723Item::Script), + Some(RunCommand::PythonStdin(contents)) => { + Pep723Stdin::parse(contents)?.map(Pep723Item::Stdin) + } _ => None, } } else if let ProjectCommand::Remove(uv_cli::RemoveArgs { @@ -226,7 +228,7 @@ async fn run(mut cli: Cli) -> Result { .. }) = &**command { - Pep723Script::read(&script).await? + Pep723Script::read(&script).await?.map(Pep723Item::Script) } else { None } @@ -237,7 +239,7 @@ async fn run(mut cli: Cli) -> Result { // If the target is a PEP 723 script, merge the metadata into the filesystem metadata. let filesystem = script .as_ref() - .map(|script| &script.metadata) + .map(Pep723Item::metadata) .and_then(|metadata| metadata.tool.as_ref()) .and_then(|tool| tool.uv.as_ref()) .map(|uv| Options::simple(uv.globals.clone(), uv.top_level.clone())) @@ -1233,7 +1235,7 @@ async fn run_project( project_command: Box, project_dir: &Path, command: Option, - script: Option, + script: Option, globals: GlobalSettings, // TODO(zanieb): Determine a better story for passing `no_config` in here no_config: bool, @@ -1465,6 +1467,12 @@ async fn run_project( .combine(Refresh::from(args.settings.upgrade.clone())), ); + // Unwrap the script. + let script = script.map(|script| match script { + Pep723Item::Script(script) => script, + Pep723Item::Stdin(_) => unreachable!("`uv remove` does not support stdin"), + }); + commands::remove( project_dir, args.locked,