Skip to content

Commit

Permalink
Support PEP 723 scripts in uv add and uv remove
Browse files Browse the repository at this point in the history
  • Loading branch information
@blueraft
blueraft committed Aug 10, 2024
1 parent 2822dde commit 6f7dd38
Show file tree
Hide file tree
Showing 11 changed files with 616 additions and 195 deletions.
7 changes: 7 additions & 0 deletions crates/uv-cli/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2462,6 +2462,10 @@ pub struct AddArgs {
help_heading = "Python options"
)]
pub python: Option<String>,

/// Specifies the Python script where the dependency will be added.
#[arg(long)]
pub script: Option<PathBuf>,
}

#[derive(Args)]
Expand Down Expand Up @@ -2509,6 +2513,9 @@ pub struct RemoveArgs {
#[arg(long, conflicts_with = "isolated")]
pub package: Option<PackageName>,

/// Specifies the Python script where the dependency will be removed.
#[arg(long)]
pub script: Option<PathBuf>,
/// The Python interpreter to use for resolving and syncing.
///
/// See `uv help python` for details on Python discovery and supported
Expand Down
177 changes: 124 additions & 53 deletions crates/uv-scripts/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,19 +19,43 @@ static FINDER: LazyLock<Finder> = LazyLock::new(|| Finder::new(b"# /// script"))
pub struct Pep723Script {
pub path: PathBuf,
pub metadata: Pep723Metadata,
pub data: String,
}

impl Pep723Script {
/// Read the PEP 723 `script` metadata from a Python file, if it exists.
///
/// See: <https://peps.python.org/pep-0723/>
pub async fn read(file: impl AsRef<Path>) -> Result<Option<Self>, Pep723Error> {
let metadata = Pep723Metadata::read(&file).await?;
Ok(metadata.map(|metadata| Self {
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()),
};

// Extract the `script` tag.
let Some((metadata, data)) = extract_script_tag(&contents)? else {
return Ok(None);
};

// Parse the metadata.
let metadata = Pep723Metadata::from_string(metadata)?;

Ok(Some(Self {
path: file.as_ref().to_path_buf(),
metadata,
data,
}))
}

/// Replace the existing metadata in the file with new metadata and write the updated content.
pub async fn replace_metadata(&self, new_metadata: &str) -> Result<(), Pep723Error> {
let new_content = format!("{}{}", serialize_metadata(new_metadata), self.data);

fs_err::tokio::write(&self.path, new_content)
.await
.map_err(std::convert::Into::into)
}
}

/// PEP 723 metadata as parsed from a `script` comment block.
Expand All @@ -43,28 +67,16 @@ pub struct Pep723Metadata {
pub dependencies: Option<Vec<pep508_rs::Requirement<VerbatimParsedUrl>>>,
pub requires_python: Option<pep440_rs::VersionSpecifiers>,
pub tool: Option<Tool>,
/// The raw unserialized document.
#[serde(skip)]
pub raw: String,
}

impl Pep723Metadata {
/// Read the PEP 723 `script` metadata from a Python file, if it exists.
///
/// See: <https://peps.python.org/pep-0723/>
pub async fn read(file: impl AsRef<Path>) -> Result<Option<Self>, Pep723Error> {
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()),
};

// Extract the `script` tag.
let Some(contents) = extract_script_tag(&contents)? else {
return Ok(None);
};

// Parse the metadata.
let metadata = toml::from_str(&contents)?;

Ok(Some(metadata))
/// Parse `Pep723Metadata` from a raw TOML string.
pub fn from_string(raw: String) -> Result<Self, toml::de::Error> {
let metadata = toml::from_str(&raw)?;
Ok(Pep723Metadata { raw, ..metadata })
}
}

Expand Down Expand Up @@ -94,34 +106,11 @@ pub enum Pep723Error {
Toml(#[from] toml::de::Error),
}

/// Read the PEP 723 `script` metadata from a Python file, if it exists.
///
/// See: <https://peps.python.org/pep-0723/>
pub async fn read_pep723_metadata(
file: impl AsRef<Path>,
) -> Result<Option<Pep723Metadata>, Pep723Error> {
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()),
};

// Extract the `script` tag.
let Some(contents) = extract_script_tag(&contents)? else {
return Ok(None);
};

// Parse the metadata.
let metadata = toml::from_str(&contents)?;

Ok(Some(metadata))
}

/// Given the contents of a Python file, extract the `script` metadata block, with leading comment
/// hashes removed.
/// hashes removed and the python script.
///
/// See: <https://peps.python.org/pep-0723/>
fn extract_script_tag(contents: &[u8]) -> Result<Option<String>, Pep723Error> {
fn extract_script_tag(contents: &[u8]) -> Result<Option<(String, String)>, Pep723Error> {
// Identify the opening pragma.
let Some(index) = FINDER.find(contents) else {
return Ok(None);
Expand Down Expand Up @@ -149,9 +138,14 @@ fn extract_script_tag(contents: &[u8]) -> Result<Option<String>, Pep723Error> {
// > second character is a space, otherwise just the first character (which means the line
// > consists of only a single #).
let mut toml = vec![];
for line in lines {

let mut python_script = vec![];

while let Some(line) = lines.next() {
// Remove the leading `#`.
let Some(line) = line.strip_prefix('#') else {
python_script.push(line);
python_script.extend(lines);
break;
};

Expand All @@ -163,11 +157,13 @@ fn extract_script_tag(contents: &[u8]) -> Result<Option<String>, Pep723Error> {

// Otherwise, the line _must_ start with ` `.
let Some(line) = line.strip_prefix(' ') else {
python_script.push(line);
python_script.extend(lines);
break;
};

toml.push(line);
}

// Find the closing `# ///`. The precedence is such that we need to identify the _last_ such
// line.
//
Expand Down Expand Up @@ -202,12 +198,36 @@ fn extract_script_tag(contents: &[u8]) -> Result<Option<String>, Pep723Error> {

// Join the lines into a single string.
let toml = toml.join("\n") + "\n";
let python_script = python_script.join("\n") + "\n";

Ok(Some((toml, python_script)))
}

/// Formats the provided metadata by prefixing each line with `#` and wrapping it with script markers.
fn serialize_metadata(metadata: &str) -> String {
let mut output = String::with_capacity(metadata.len() + 2);

output.push_str("# /// script\n");

Ok(Some(toml))
for line in metadata.lines() {
if line.is_empty() {
output.push('\n');
} else {
output.push_str("# ");
output.push_str(line);
output.push('\n');
}
}

output.push_str("# ///\n");

output
}

#[cfg(test)]
mod tests {
use crate::serialize_metadata;

#[test]
fn missing_space() {
let contents = indoc::indoc! {r"
Expand Down Expand Up @@ -269,21 +289,37 @@ mod tests {
# 'rich',
# ]
# ///
import requests
from rich.pretty import pprint
resp = requests.get('https://peps.python.org/api/peps.json')
data = resp.json()
"};

let expected = indoc::indoc! {r"
let expected_metadata = indoc::indoc! {r"
requires-python = '>=3.11'
dependencies = [
'requests<3',
'rich',
]
"};

let expected_data = indoc::indoc! {r"
import requests
from rich.pretty import pprint
resp = requests.get('https://peps.python.org/api/peps.json')
data = resp.json()
"};

let actual = super::extract_script_tag(contents.as_bytes())
.unwrap()
.unwrap();

assert_eq!(actual, expected);
assert_eq!(actual.0, expected_metadata);
assert_eq!(actual.1, expected_data);
}

#[test]
Expand Down Expand Up @@ -312,7 +348,8 @@ mod tests {

let actual = super::extract_script_tag(contents.as_bytes())
.unwrap()
.unwrap();
.unwrap()
.0;

assert_eq!(actual, expected);
}
Expand Down Expand Up @@ -341,8 +378,42 @@ mod tests {

let actual = super::extract_script_tag(contents.as_bytes())
.unwrap()
.unwrap();
.unwrap()
.0;

assert_eq!(actual, expected);
}

#[test]
fn test_serialize_metadata_formatting() {
let metadata = indoc::indoc! {r"
requires-python = '>=3.11'
dependencies = [
'requests<3',
'rich',
]
"};

let expected_output = indoc::indoc! {r"
# /// script
# requires-python = '>=3.11'
# dependencies = [
# 'requests<3',
# 'rich',
# ]
# ///
"};

let result = serialize_metadata(metadata);
assert_eq!(result, expected_output);
}

#[test]
fn test_serialize_metadata_empty() {
let metadata = "";
let expected_output = "# /// script\n# ///\n";

let result = serialize_metadata(metadata);
assert_eq!(result, expected_output);
}
}
Loading

0 comments on commit 6f7dd38

Please sign in to comment.