diff --git a/crates/uv-workspace/src/pyproject_mut.rs b/crates/uv-workspace/src/pyproject_mut.rs index 3befdd5ac..696aa442f 100644 --- a/crates/uv-workspace/src/pyproject_mut.rs +++ b/crates/uv-workspace/src/pyproject_mut.rs @@ -49,6 +49,20 @@ pub enum ArrayEdit { Add(usize), } +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +enum CommentType { + /// A comment that appears on its own line. + OwnLine, + /// A comment that appears at the end of a line. + EndOfLine, +} + +#[derive(Debug, Clone)] +struct Comment { + text: String, + comment_type: CommentType, +} + impl ArrayEdit { pub fn index(&self) -> usize { match self { @@ -694,7 +708,12 @@ pub fn add_dependency( }; let index = index.unwrap_or(deps.len()); - deps.insert(index, req_string); + let mut value = Value::from(req_string.as_str()); + let decor = value.decor_mut(); + decor.set_prefix(deps.trailing().clone()); + deps.set_trailing(""); + + deps.insert_formatted(index, value); // `reformat_array_multiline` uses the indentation of the first dependency entry. // Therefore, we retrieve the indentation of the first dependency entry and apply it to // the new entry. Note that it is only necessary if the newly added dependency is going @@ -829,14 +848,38 @@ fn try_parse_requirement(req: &str) -> Option { /// Reformats a TOML array to multi line while trying to preserve all comments /// and move them around. This also formats the array to have a trailing comma. fn reformat_array_multiline(deps: &mut Array) { - fn find_comments(s: Option<&RawString>) -> impl Iterator { - s.and_then(|x| x.as_str()) + fn find_comments(s: Option<&RawString>) -> Box + '_> { + let iter = s + .and_then(|x| x.as_str()) .unwrap_or("") .lines() - .filter_map(|line| { - let line = line.trim(); - line.starts_with('#').then_some(line) - }) + .scan( + (false, false), + |(prev_line_was_empty, prev_line_was_comment), line| { + let trimmed_line = line.trim(); + if let Some(index) = trimmed_line.find('#') { + let comment_text = trimmed_line[index..].trim().to_string(); + let comment_type = if (*prev_line_was_empty) || (*prev_line_was_comment) { + CommentType::OwnLine + } else { + CommentType::EndOfLine + }; + *prev_line_was_empty = trimmed_line.is_empty(); + *prev_line_was_comment = true; + Some(Some(Comment { + text: comment_text, + comment_type, + })) + } else { + *prev_line_was_empty = trimmed_line.is_empty(); + *prev_line_was_comment = false; + Some(None) + } + }, + ) + .flatten(); + + Box::new(iter) } let mut indentation_prefix = None; @@ -866,8 +909,15 @@ fn reformat_array_multiline(deps: &mut Array) { let indentation_prefix_str = format!("\n{}", indentation_prefix.as_ref().unwrap()); for comment in find_comments(decor.prefix()).chain(find_comments(decor.suffix())) { - prefix.push_str(&indentation_prefix_str); - prefix.push_str(comment); + match comment.comment_type { + CommentType::OwnLine => { + prefix.push_str(&indentation_prefix_str); + } + CommentType::EndOfLine => { + prefix.push(' '); + } + } + prefix.push_str(&comment.text); } prefix.push_str(&indentation_prefix_str); decor.set_prefix(prefix); @@ -880,7 +930,7 @@ fn reformat_array_multiline(deps: &mut Array) { if comments.peek().is_some() { for comment in comments { rv.push_str("\n "); - rv.push_str(comment); + rv.push_str(&comment.text); } } if !rv.is_empty() || !deps.is_empty() { diff --git a/crates/uv/tests/it/edit.rs b/crates/uv/tests/it/edit.rs index ffbdf0cc0..94607c904 100644 --- a/crates/uv/tests/it/edit.rs +++ b/crates/uv/tests/it/edit.rs @@ -6269,3 +6269,137 @@ fn add_self() -> Result<()> { Ok(()) } + +#[test] +fn add_preserves_end_of_line_comments() -> Result<()> { + let context = TestContext::new("3.12"); + + let pyproject_toml = context.temp_dir.child("pyproject.toml"); + pyproject_toml.write_str(indoc! {r#" + [project] + name = "project" + version = "0.1.0" + requires-python = ">=3.12" + dependencies = [ + # comment 0 + "anyio==3.7.0", # comment 1 + # comment 2 + ] + + [build-system] + requires = ["setuptools>=42"] + build-backend = "setuptools.build_meta" + "#})?; + + uv_snapshot!(context.filters(), context.add().arg("requests==2.31.0"), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Resolved 8 packages in [TIME] + Prepared 8 packages in [TIME] + Installed 8 packages in [TIME] + + anyio==3.7.0 + + certifi==2024.2.2 + + charset-normalizer==3.3.2 + + idna==3.6 + + project==0.1.0 (from file://[TEMP_DIR]/) + + requests==2.31.0 + + sniffio==1.3.1 + + urllib3==2.2.1 + "###); + + let pyproject_toml = context.read("pyproject.toml"); + + insta::with_settings!({ + filters => context.filters(), + }, { + assert_snapshot!( + pyproject_toml, @r###" + [project] + name = "project" + version = "0.1.0" + requires-python = ">=3.12" + dependencies = [ + # comment 0 + "anyio==3.7.0", # comment 1 + # comment 2 + "requests==2.31.0", + ] + + [build-system] + requires = ["setuptools>=42"] + build-backend = "setuptools.build_meta" + "### + ); + }); + Ok(()) +} + +#[test] +fn add_preserves_open_bracket_comment() -> Result<()> { + let context = TestContext::new("3.12"); + + let pyproject_toml = context.temp_dir.child("pyproject.toml"); + pyproject_toml.write_str(indoc! {r#" + [project] + name = "project" + version = "0.1.0" + requires-python = ">=3.12" + dependencies = [ # comment 0 + # comment 1 + "anyio==3.7.0", # comment 2 + # comment 3 + ] + + [build-system] + requires = ["setuptools>=42"] + build-backend = "setuptools.build_meta" + "#})?; + + uv_snapshot!(context.filters(), context.add().arg("requests==2.31.0"), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Resolved 8 packages in [TIME] + Prepared 8 packages in [TIME] + Installed 8 packages in [TIME] + + anyio==3.7.0 + + certifi==2024.2.2 + + charset-normalizer==3.3.2 + + idna==3.6 + + project==0.1.0 (from file://[TEMP_DIR]/) + + requests==2.31.0 + + sniffio==1.3.1 + + urllib3==2.2.1 + "###); + + let pyproject_toml = context.read("pyproject.toml"); + + insta::with_settings!({ + filters => context.filters(), + }, { + assert_snapshot!( + pyproject_toml, @r###" + [project] + name = "project" + version = "0.1.0" + requires-python = ">=3.12" + dependencies = [ # comment 0 + # comment 1 + "anyio==3.7.0", # comment 2 + # comment 3 + "requests==2.31.0", + ] + + [build-system] + requires = ["setuptools>=42"] + build-backend = "setuptools.build_meta" + "### + ); + }); + Ok(()) +}