diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/while.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/while.py new file mode 100644 index 0000000000..79c0d09c8c --- /dev/null +++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/while.py @@ -0,0 +1,30 @@ +while 34: # trailing test comment + pass # trailing last statement comment + + # trailing while body comment + +# leading else comment + +else: # trailing else comment + pass + + # trailing else body comment + + +while aVeryLongConditionThatSpillsOverToTheNextLineBecauseItIsExtremelyLongAndGoesOnAndOnAndOnAndOnAndOnAndOnAndOnAndOnAndOn: # trailing comment + pass + +else: + ... + +while ( + some_condition(unformatted, args) and anotherCondition or aThirdCondition +): # comment + print("Do something") + + +while ( + some_condition(unformatted, args) # trailing some condition + and anotherCondition or aThirdCondition # trailing third condition +): # comment + print("Do something") diff --git a/crates/ruff_python_formatter/src/builders.rs b/crates/ruff_python_formatter/src/builders.rs index 62d7317f36..45a9ee248e 100644 --- a/crates/ruff_python_formatter/src/builders.rs +++ b/crates/ruff_python_formatter/src/builders.rs @@ -11,7 +11,7 @@ pub(crate) trait PyFormatterExtensions<'ast, 'buf> { /// empty lines between any two nodes. Separates any two nodes by at least a hard line break. /// /// * [`NodeLevel::Module`]: Up to two empty lines - /// * [`NodeLevel::Statement`]: Up to one empty line + /// * [`NodeLevel::CompoundStatement`]: Up to one empty line /// * [`NodeLevel::Parenthesized`]: No empty lines fn join_nodes<'fmt>(&'fmt mut self, level: NodeLevel) -> JoinNodesBuilder<'fmt, 'ast, 'buf>; } @@ -53,10 +53,12 @@ impl<'fmt, 'ast, 'buf> JoinNodesBuilder<'fmt, 'ast, 'buf> { 2 => empty_line().fmt(f), _ => write!(f, [empty_line(), empty_line()]), }, - NodeLevel::Statement => match lines_before(f.context().contents(), node.start()) { - 0 | 1 => hard_line_break().fmt(f), - _ => empty_line().fmt(f), - }, + NodeLevel::CompoundStatement => { + match lines_before(f.context().contents(), node.start()) { + 0 | 1 => hard_line_break().fmt(f), + _ => empty_line().fmt(f), + } + } NodeLevel::Parenthesized => hard_line_break().fmt(f), }); @@ -180,7 +182,7 @@ no_leading_newline = 30"# // Should keep at most one empty level #[test] fn ranged_builder_statement_level() { - let printed = format_ranged(NodeLevel::Statement); + let printed = format_ranged(NodeLevel::CompoundStatement); assert_eq!( &printed, diff --git a/crates/ruff_python_formatter/src/comments/format.rs b/crates/ruff_python_formatter/src/comments/format.rs index 1eb2f44184..286905ce85 100644 --- a/crates/ruff_python_formatter/src/comments/format.rs +++ b/crates/ruff_python_formatter/src/comments/format.rs @@ -6,27 +6,37 @@ use ruff_formatter::{format_args, write, FormatError, SourceCode}; use ruff_python_ast::node::AnyNodeRef; use ruff_python_ast::prelude::AstNode; use ruff_text_size::{TextLen, TextRange, TextSize}; +use rustpython_parser::ast::Ranged; /// Formats the leading comments of a node. -pub(crate) fn leading_comments(node: &T) -> FormatLeadingComments +pub(crate) fn leading_node_comments(node: &T) -> FormatLeadingComments where T: AstNode, { - FormatLeadingComments { - node: node.as_any_node_ref(), - } + FormatLeadingComments::Node(node.as_any_node_ref()) +} + +/// Formats the passed comments as leading comments +pub(crate) const fn leading_comments(comments: &[SourceComment]) -> FormatLeadingComments { + FormatLeadingComments::Comments(comments) } #[derive(Copy, Clone, Debug)] -pub(crate) struct FormatLeadingComments<'a> { - node: AnyNodeRef<'a>, +pub(crate) enum FormatLeadingComments<'a> { + Node(AnyNodeRef<'a>), + Comments(&'a [SourceComment]), } impl Format> for FormatLeadingComments<'_> { fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> { let comments = f.context().comments().clone(); - for comment in comments.leading_comments(self.node) { + let leading_comments = match self { + FormatLeadingComments::Node(node) => comments.leading_comments(*node), + FormatLeadingComments::Comments(comments) => comments, + }; + + for comment in leading_comments { let slice = comment.slice(); let lines_after_comment = lines_after(f.context().contents(), slice.end()); @@ -42,32 +52,88 @@ impl Format> for FormatLeadingComments<'_> { } } -/// Formats the trailing comments of `node` -pub(crate) fn trailing_comments(node: &T) -> FormatTrailingComments +/// Formats the leading `comments` of an alternate branch and ensures that it preserves the right +/// number of empty lines before. The `last_node` is the last node of the preceding body. +/// +/// For example, `last_node` is the last statement in the if body when formatting the leading +/// comments of the `else` branch. +pub(crate) fn leading_alternate_branch_comments<'a, T>( + comments: &'a [SourceComment], + last_node: Option, +) -> FormatLeadingAlternateBranchComments<'a> where - T: AstNode, + T: Into>, { - FormatTrailingComments { - node: node.as_any_node_ref(), + FormatLeadingAlternateBranchComments { + comments, + last_node: last_node.map(std::convert::Into::into), } } -pub(crate) struct FormatTrailingComments<'a> { - node: AnyNodeRef<'a>, +pub(crate) struct FormatLeadingAlternateBranchComments<'a> { + comments: &'a [SourceComment], + last_node: Option>, +} + +impl Format> for FormatLeadingAlternateBranchComments<'_> { + fn fmt(&self, f: &mut Formatter>) -> FormatResult<()> { + if let Some(first_leading) = self.comments.first() { + // Leading comments only preserves the lines after the comment but not before. + // Insert the necessary lines. + if lines_before(f.context().contents(), first_leading.slice().start()) > 1 { + write!(f, [empty_line()])?; + } + + write!(f, [leading_comments(self.comments)])?; + } else if let Some(last_preceding) = self.last_node { + // The leading comments formatting ensures that it preserves the right amount of lines after + // We need to take care of this ourselves, if there's no leading `else` comment. + if lines_after(f.context().contents(), last_preceding.end()) > 1 { + write!(f, [empty_line()])?; + } + } + + Ok(()) + } +} + +/// Formats the trailing comments of `node` +pub(crate) fn trailing_node_comments(node: &T) -> FormatTrailingComments +where + T: AstNode, +{ + FormatTrailingComments::Node(node.as_any_node_ref()) +} + +/// Formats the passed comments as trailing comments +pub(crate) fn trailing_comments(comments: &[SourceComment]) -> FormatTrailingComments { + FormatTrailingComments::Comments(comments) +} + +pub(crate) enum FormatTrailingComments<'a> { + Node(AnyNodeRef<'a>), + Comments(&'a [SourceComment]), } impl Format> for FormatTrailingComments<'_> { fn fmt(&self, f: &mut Formatter>) -> FormatResult<()> { let comments = f.context().comments().clone(); - let mut has_empty_lines_before = false; - for trailing in comments.trailing_comments(self.node) { + let trailing_comments = match self { + FormatTrailingComments::Node(node) => comments.trailing_comments(*node), + FormatTrailingComments::Comments(comments) => comments, + }; + + let mut has_trailing_own_line_comment = false; + + for trailing in trailing_comments { let slice = trailing.slice(); - let lines_before_comment = lines_before(f.context().contents(), slice.start()); - has_empty_lines_before |= lines_before_comment > 0; + has_trailing_own_line_comment |= trailing.position().is_own_line(); + + if has_trailing_own_line_comment { + let lines_before_comment = lines_before(f.context().contents(), slice.start()); - if has_empty_lines_before { // A trailing comment at the end of a body or list // ```python // def test(): @@ -105,7 +171,7 @@ impl Format> for FormatTrailingComments<'_> { } /// Formats the dangling comments of `node`. -pub(crate) fn dangling_comments(node: &T) -> FormatDanglingComments +pub(crate) fn dangling_node_comments(node: &T) -> FormatDanglingComments where T: AstNode, { @@ -229,7 +295,7 @@ impl Format> for FormatEmptyLines { _ => write!(f, [empty_line(), empty_line()]), }, - NodeLevel::Statement => match self.lines { + NodeLevel::CompoundStatement => match self.lines { 0 | 1 => write!(f, [hard_line_break()]), _ => write!(f, [empty_line()]), }, diff --git a/crates/ruff_python_formatter/src/comments/mod.rs b/crates/ruff_python_formatter/src/comments/mod.rs index 57206e3fc4..9476ccda45 100644 --- a/crates/ruff_python_formatter/src/comments/mod.rs +++ b/crates/ruff_python_formatter/src/comments/mod.rs @@ -103,7 +103,10 @@ use crate::comments::debug::{DebugComment, DebugComments}; use crate::comments::map::MultiMap; use crate::comments::node_key::NodeRefEqualityKey; use crate::comments::visitor::CommentsVisitor; -pub(crate) use format::{dangling_comments, leading_comments, trailing_comments}; +pub(crate) use format::{ + dangling_node_comments, leading_alternate_branch_comments, leading_node_comments, + trailing_comments, trailing_node_comments, +}; use ruff_formatter::{SourceCode, SourceCodeSlice}; use ruff_python_ast::node::AnyNodeRef; use ruff_python_ast::source_code::CommentRanges; @@ -121,8 +124,6 @@ pub(crate) struct SourceComment { position: CommentTextPosition, } -#[allow(unused)] -// TODO(micha): Remove after using the new comments infrastructure in the formatter. impl SourceComment { /// Returns the location of the comment in the original source code. /// Allows retrieving the text of the comment. @@ -184,8 +185,6 @@ pub(crate) enum CommentTextPosition { OwnLine, } -#[allow(unused)] -// TODO(micha): Remove after using the new comments infrastructure in the formatter. impl CommentTextPosition { pub(crate) const fn is_own_line(self) -> bool { matches!(self, CommentTextPosition::OwnLine) diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index d4b007e775..062cd17a8d 100644 --- a/crates/ruff_python_formatter/src/comments/placement.rs +++ b/crates/ruff_python_formatter/src/comments/placement.rs @@ -5,7 +5,7 @@ use crate::trivia::find_first_non_trivia_character_in_range; use ruff_python_ast::node::AnyNodeRef; use ruff_python_ast::source_code::Locator; use ruff_python_ast::whitespace; -use ruff_text_size::{TextRange, TextSize}; +use ruff_text_size::{TextLen, TextRange, TextSize}; use rustpython_parser::ast::Ranged; use std::cmp::Ordering; @@ -20,6 +20,7 @@ pub(super) fn place_comment<'a>( .or_else(|comment| handle_in_between_bodies_end_of_line_comment(comment, locator)) .or_else(|comment| handle_trailing_body_comment(comment, locator)) .or_else(handle_trailing_end_of_line_body_comment) + .or_else(|comment| handle_trailing_end_of_line_condition_comment(comment, locator)) .or_else(|comment| handle_positional_only_arguments_separator_comment(comment, locator)) .or_else(|comment| { handle_trailing_binary_expression_left_or_operator_comment(comment, locator) @@ -471,6 +472,91 @@ fn handle_trailing_end_of_line_body_comment(comment: DecoratedComment<'_>) -> Co } } +/// Handles end of line comments after the `:` of a condition +/// +/// ```python +/// while True: # comment +/// pass +/// ``` +/// +/// It attaches the comment as dangling comment to the enclosing `while` statement. +fn handle_trailing_end_of_line_condition_comment<'a>( + comment: DecoratedComment<'a>, + locator: &Locator, +) -> CommentPlacement<'a> { + use ruff_python_ast::prelude::*; + + // Must be an end of line comment + if comment.text_position().is_own_line() { + return CommentPlacement::Default(comment); + } + + // Must be between the condition expression and the first body element + let (Some(preceding), Some(following)) = (comment.preceding_node(), comment.following_node()) else { + return CommentPlacement::Default(comment); + }; + + let expression_before_colon = match comment.enclosing_node() { + AnyNodeRef::StmtIf(StmtIf { test: expr, .. }) + | AnyNodeRef::StmtWhile(StmtWhile { test: expr, .. }) + | AnyNodeRef::StmtFor(StmtFor { iter: expr, .. }) + | AnyNodeRef::StmtAsyncFor(StmtAsyncFor { iter: expr, .. }) => { + Some(AnyNodeRef::from(expr.as_ref())) + } + + AnyNodeRef::StmtWith(StmtWith { items, .. }) + | AnyNodeRef::StmtAsyncWith(StmtAsyncWith { items, .. }) => { + items.last().map(AnyNodeRef::from) + } + _ => None, + }; + + let Some(last_before_colon) = expression_before_colon else { + return CommentPlacement::Default(comment); + }; + + // If the preceding is the node before the `colon` + // `while true:` The node before the `colon` is the `true` constant. + if preceding.ptr_eq(last_before_colon) { + let mut start = preceding.end(); + while let Some((offset, c)) = find_first_non_trivia_character_in_range( + locator.contents(), + TextRange::new(start, following.start()), + ) { + match c { + ':' => { + if comment.slice().start() > offset { + // Comment comes after the colon + // ```python + // while a: # comment + // ... + // ``` + return CommentPlacement::dangling(comment.enclosing_node(), comment); + } + + // Comment comes before the colon + // ```python + // while ( + // a # comment + // ): + // ... + // ``` + break; + } + ')' => { + // Skip over any closing parentheses + start = offset + ')'.text_len(); + } + _ => { + unreachable!("Only ')' or ':' should follow the condition") + } + } + } + } + + CommentPlacement::Default(comment) +} + /// Attaches comments for the positional-only arguments separator `/` as trailing comments to the /// enclosing [`Arguments`] node. /// diff --git a/crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__base_test.snap b/crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__base_test.snap index 51d2b14f09..ef24b1dbef 100644 --- a/crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__base_test.snap +++ b/crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__base_test.snap @@ -19,19 +19,19 @@ expression: comments.debug(test_case.source_code) "trailing": [], }, Node { - kind: ExprCompare, - range: 51..57, - source: `x == y`, + kind: StmtIf, + range: 48..212, + source: `if x == y: # if statement e...ne comment⏎`, }: { "leading": [], - "dangling": [], - "trailing": [ + "dangling": [ SourceComment { text: "# if statement end of line comment", position: EndOfLine, formatted: false, }, ], + "trailing": [], }, Node { kind: StmtIf, diff --git a/crates/ruff_python_formatter/src/context.rs b/crates/ruff_python_formatter/src/context.rs index 21d03fec72..bef99c4514 100644 --- a/crates/ruff_python_formatter/src/context.rs +++ b/crates/ruff_python_formatter/src/context.rs @@ -25,7 +25,6 @@ impl<'a> PyFormatContext<'a> { } } - #[allow(unused)] pub(crate) fn contents(&self) -> &'a str { self.contents } @@ -35,7 +34,6 @@ impl<'a> PyFormatContext<'a> { Locator::new(self.contents) } - #[allow(unused)] pub(crate) fn set_node_level(&mut self, level: NodeLevel) { self.node_level = level; } @@ -44,7 +42,6 @@ impl<'a> PyFormatContext<'a> { self.node_level } - #[allow(unused)] pub(crate) fn comments(&self) -> &Comments<'a> { &self.comments } @@ -80,11 +77,10 @@ pub(crate) enum NodeLevel { #[default] TopLevel, - /// Formatting nodes that are enclosed by a statement. - #[allow(unused)] - Statement, + /// Formatting the body statements of a [compound statement](https://docs.python.org/3/reference/compound_stmts.html#compound-statements) + /// (`if`, `while`, `match`, etc.). + CompoundStatement, /// Formatting nodes that are enclosed in a parenthesized expression. - #[allow(unused)] Parenthesized, } diff --git a/crates/ruff_python_formatter/src/expression/maybe_parenthesize.rs b/crates/ruff_python_formatter/src/expression/maybe_parenthesize.rs new file mode 100644 index 0000000000..aa7b4f5702 --- /dev/null +++ b/crates/ruff_python_formatter/src/expression/maybe_parenthesize.rs @@ -0,0 +1,53 @@ +use crate::context::NodeLevel; +use crate::prelude::*; +use ruff_formatter::{format_args, write}; +use rustpython_parser::ast::Expr; + +/// Formats the passed expression. Adds parentheses if the expression doesn't fit on a line. +pub(crate) const fn maybe_parenthesize(expression: &Expr) -> MaybeParenthesize { + MaybeParenthesize { expression } +} + +pub(crate) struct MaybeParenthesize<'a> { + expression: &'a Expr, +} + +impl Format> for MaybeParenthesize<'_> { + fn fmt(&self, f: &mut Formatter>) -> FormatResult<()> { + let saved_level = f.context().node_level(); + f.context_mut().set_node_level(NodeLevel::Parenthesized); + + let result = if needs_parentheses(self.expression) { + write!( + f, + [group(&format_args![ + if_group_breaks(&text("(")), + soft_block_indent(&self.expression.format()), + if_group_breaks(&text(")")) + ])] + ) + } else { + // Don't add parentheses around expressions that have parentheses on their own (e.g. list, dict, tuple, call expression) + self.expression.format().fmt(f) + }; + + f.context_mut().set_node_level(saved_level); + + result + } +} + +const fn needs_parentheses(expr: &Expr) -> bool { + !matches!( + expr, + Expr::Tuple(_) + | Expr::List(_) + | Expr::Set(_) + | Expr::Dict(_) + | Expr::ListComp(_) + | Expr::SetComp(_) + | Expr::DictComp(_) + | Expr::GeneratorExp(_) + | Expr::Call(_) + ) +} diff --git a/crates/ruff_python_formatter/src/expression/mod.rs b/crates/ruff_python_formatter/src/expression/mod.rs index 362b1daa86..039c104d55 100644 --- a/crates/ruff_python_formatter/src/expression/mod.rs +++ b/crates/ruff_python_formatter/src/expression/mod.rs @@ -29,6 +29,7 @@ pub(crate) mod expr_tuple; pub(crate) mod expr_unary_op; pub(crate) mod expr_yield; pub(crate) mod expr_yield_from; +pub(crate) mod maybe_parenthesize; #[derive(Default)] pub struct FormatExpr; diff --git a/crates/ruff_python_formatter/src/lib.rs b/crates/ruff_python_formatter/src/lib.rs index 49cc9fadc9..dcbf68acc9 100644 --- a/crates/ruff_python_formatter/src/lib.rs +++ b/crates/ruff_python_formatter/src/lib.rs @@ -14,7 +14,9 @@ use ruff_formatter::{ use ruff_python_ast::node::AstNode; use ruff_python_ast::source_code::{CommentRanges, CommentRangesBuilder, Locator}; -use crate::comments::{dangling_comments, leading_comments, trailing_comments, Comments}; +use crate::comments::{ + dangling_node_comments, leading_node_comments, trailing_node_comments, Comments, +}; use crate::context::PyFormatContext; pub(crate) mod builders; @@ -64,7 +66,7 @@ where /// You may want to override this method if you want to manually handle the formatting of comments /// inside of the `fmt_fields` method or customize the formatting of the leading comments. fn fmt_leading_comments(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> { - leading_comments(node).fmt(f) + leading_node_comments(node).fmt(f) } /// Formats the [dangling comments](comments#dangling-comments) of the node. @@ -75,7 +77,7 @@ where /// /// A node can have dangling comments if all its children are tokens or if all node childrens are optional. fn fmt_dangling_comments(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> { - dangling_comments(node).fmt(f) + dangling_node_comments(node).fmt(f) } /// Formats the [trailing comments](comments#trailing-comments) of the node. @@ -83,7 +85,7 @@ where /// You may want to override this method if you want to manually handle the formatting of comments /// inside of the `fmt_fields` method or customize the formatting of the trailing comments. fn fmt_trailing_comments(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> { - trailing_comments(node).fmt(f) + trailing_node_comments(node).fmt(f) } } @@ -285,14 +287,58 @@ Formatted twice: Ok(()) } + #[fixture(pattern = "resources/test/fixtures/ruff/**/*.py")] + #[test] + fn ruff_test(input_path: &Path) -> Result<()> { + let content = fs::read_to_string(input_path)?; + + let printed = format_module(&content)?; + let formatted_code = printed.as_code(); + + let reformatted = + format_module(formatted_code).expect("Expected formatted code to be valid syntax"); + + if reformatted.as_code() != formatted_code { + let diff = TextDiff::from_lines(formatted_code, reformatted.as_code()) + .unified_diff() + .header("Formatted once", "Formatted twice") + .to_string(); + panic!( + r#"Reformatting the formatted code a second time resulted in formatting changes. +{diff} + +Formatted once: +{formatted_code} + +Formatted twice: +{}"#, + reformatted.as_code() + ); + } + + let snapshot = format!( + r#"## Input +{} + +## Output +{}"#, + CodeFrame::new("py", &content), + CodeFrame::new("py", formatted_code) + ); + assert_snapshot!(snapshot); + + Ok(()) + } + /// Use this test to debug the formatting of some snipped #[ignore] #[test] fn quick_test() { let src = r#" -{ - k: v for k, v in a_very_long_variable_name_that_exceeds_the_line_length_by_far_keep_going -} +while True: + if something.changed: + do.stuff() # trailing comment +other "#; // Tokenize once let mut tokens = Vec::new(); @@ -320,10 +366,10 @@ Formatted twice: assert_eq!( printed.as_code(), - r#"{ - k: v - for k, v in a_very_long_variable_name_that_exceeds_the_line_length_by_far_keep_going -}"# + r#"while True: + if something.changed: + do.stuff() # trailing comment +"# ); } diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments5_py.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments5_py.snap index 7087a8c711..dcace90ac2 100644 --- a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments5_py.snap +++ b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments5_py.snap @@ -86,20 +86,20 @@ if __name__ == "__main__": ```diff --- Black +++ Ruff -@@ -1,11 +1,6 @@ +@@ -1,11 +1,9 @@ while True: if something.changed: - do.stuff() # trailing comment - # Comment belongs to the `if` block. -- # This one belongs to the `while` block. -- -- # Should this one, too? I guess so. -- + do.stuff() + # This one belongs to the `while` block. + + # Should this one, too? I guess so. +- # This one is properly standalone now. for i in range(100): -@@ -15,27 +10,18 @@ +@@ -15,27 +13,18 @@ # then we do this print(i) @@ -127,7 +127,7 @@ if __name__ == "__main__": # SECTION COMMENT -@@ -47,8 +33,6 @@ +@@ -47,8 +36,6 @@ @deco3 def decorated1(): ... @@ -136,7 +136,7 @@ if __name__ == "__main__": # leading 1 @deco1 # leading 2 -@@ -56,18 +40,12 @@ +@@ -56,18 +43,12 @@ # leading function comment def decorated1(): ... @@ -163,6 +163,9 @@ if __name__ == "__main__": while True: if something.changed: do.stuff() + # This one belongs to the `while` block. + + # Should this one, too? I guess so. # This one is properly standalone now. for i in range(100): diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__expression_py.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__expression_py.snap index cf47dd1cbe..8dbc397d78 100644 --- a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__expression_py.snap +++ b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__expression_py.snap @@ -511,7 +511,7 @@ last_call() Ø = set() authors.łukasz.say_thanks() mapping = { -@@ -233,138 +170,83 @@ +@@ -233,138 +170,84 @@ C: 0.1 * (10.0 / 12), D: 0.1 * (10.0 / 12), } @@ -550,15 +550,6 @@ last_call() - ... -for j in 1 + (2 + 3): - ... --while this and that: -- ... --for ( -- addr_family, -- addr_type, -- addr_proto, -- addr_canonname, -- addr_sockaddr, --) in socket.getaddrinfo("google.com", "http"): +print(* lambda x: x) +assert(not Test),("Short message") +assert this is ComplexTest and not requirements.fit_in_a_single_line(force=False), "Short message" @@ -568,7 +559,15 @@ last_call() +for z in (i for i in (1, 2, 3)): ... +for i in (call()): ... +for j in (1 + (2 + 3)): ... -+while(this and that): ... + while this and that: + ... +-for ( +- addr_family, +- addr_type, +- addr_proto, +- addr_canonname, +- addr_sockaddr, +-) in socket.getaddrinfo("google.com", "http"): +for addr_family, addr_type, addr_proto, addr_canonname, addr_sockaddr in socket.getaddrinfo('google.com', 'http'): pass -a = ( @@ -879,7 +878,8 @@ for y in (): ... for z in (i for i in (1, 2, 3)): ... for i in (call()): ... for j in (1 + (2 + 3)): ... -while(this and that): ... +while this and that: + ... for addr_family, addr_type, addr_proto, addr_canonname, addr_sockaddr in socket.getaddrinfo('google.com', 'http'): pass a = aaaa.bbbb.cccc.dddd.eeee.ffff.gggg.hhhh.iiii.jjjj.kkkk.llll.mmmm.nnnn.oooo.pppp in qqqq.rrrr.ssss.tttt.uuuu.vvvv.xxxx.yyyy.zzzz diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__fmtskip8_py.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__fmtskip8_py.snap index 358018b4e0..0c9e4e258b 100644 --- a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__fmtskip8_py.snap +++ b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__fmtskip8_py.snap @@ -100,20 +100,24 @@ async def test_async_with(): # Make sure a leading comment is not removed. if unformatted_call( args ): # fmt: skip print("First branch") -@@ -30,33 +23,21 @@ +@@ -29,34 +22,22 @@ + elif another_unformatted_call( args ): # fmt: skip print("Second branch") else : # fmt: skip - print("Last branch") +- print("Last branch") - - - while some_condition( unformatted, args ): # fmt: skip +-while some_condition( unformatted, args ): # fmt: skip ++ print("Last branch") # fmt: skip ++while some_condition( unformatted, args ): # fmt: skip print("Do something") - - for i in some_iter( unformatted, args ): # fmt: skip - print("Do something") +- print("Do something") - - ++ print("Do something") # fmt: skip async def test_async_for(): async for i in some_async_iter( unformatted, args ): # fmt: skip print("Do something") @@ -128,9 +132,10 @@ async def test_async_with(): - - with give_me_context( unformatted, args ): # fmt: skip - print("Do something") +- print("Do something") - - ++ print("Do something") # fmt: skip async def test_async_with(): async with give_me_async_context( unformatted, args ): # fmt: skip print("Do something") @@ -163,11 +168,11 @@ if unformatted_call( args ): # fmt: skip elif another_unformatted_call( args ): # fmt: skip print("Second branch") else : # fmt: skip - print("Last branch") -while some_condition( unformatted, args ): # fmt: skip + print("Last branch") # fmt: skip +while some_condition( unformatted, args ): # fmt: skip print("Do something") for i in some_iter( unformatted, args ): # fmt: skip - print("Do something") + print("Do something") # fmt: skip async def test_async_for(): async for i in some_async_iter( unformatted, args ): # fmt: skip print("Do something") @@ -178,7 +183,7 @@ except UnformattedError as ex: # fmt: skip finally : # fmt: skip finally_call() with give_me_context( unformatted, args ): # fmt: skip - print("Do something") + print("Do something") # fmt: skip async def test_async_with(): async with give_me_async_context( unformatted, args ): # fmt: skip print("Do something") diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__remove_newline_after_code_block_open_py.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__remove_newline_after_code_block_open_py.snap index bb3d5adbec..fcfe41d5a6 100644 --- a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__remove_newline_after_code_block_open_py.snap +++ b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__remove_newline_after_code_block_open_py.snap @@ -121,30 +121,30 @@ with open("/path/to/file.txt", mode="r") as read_file: ```diff --- Black +++ Ruff -@@ -1,78 +1,74 @@ +@@ -1,78 +1,68 @@ import random -- -- - def foo1(): -- print("The newline above me should be deleted!") -- ++def foo1(): + print("The newline above me should be deleted!") - def foo2(): -- print("All the newlines above me should be deleted!") ++def foo2(): + +-def foo1(): +- print("The newline above me should be deleted!") -+ -+ print("All the newlines above me should be deleted!") +-def foo2(): + print("All the newlines above me should be deleted!") +- +- def foo3(): + print("No newline above me!") print("There is a newline above me, and that's OK!") -+def foo4(): - - --def foo4(): +- + def foo4(): ++ # There is a comment here print("The newline above me should not be deleted!") @@ -154,23 +154,23 @@ with open("/path/to/file.txt", mode="r") as read_file: def bar(self): + print("The newline above me should be deleted!") -- -- - for i in range(5): -- print(f"{i}) The line above me should be removed!") -- ++for i in range(5): + print(f"{i}) The line above me should be removed!") ++for i in range(5): + +-for i in range(5): +- print(f"{i}) The line above me should be removed!") + + ++ print(f"{i}) The lines above me should be removed!") for i in range(5): - print(f"{i}) The lines above me should be removed!") ++ for j in range(7): -+ -+ print(f"{i}) The lines above me should be removed!") - for i in range(5): -+ - for j in range(7): -+ +-for i in range(5): +- for j in range(7): print(f"{i}) The lines above me should be removed!") +if random.randint(0, 3) == 0: @@ -189,38 +189,33 @@ with open("/path/to/file.txt", mode="r") as read_file: if random.uniform(0, 1) > 0.5: print("Two lines above me are about to be removed!") - -+while True: - -+ print("The newline above me should be deleted!") +- + while True: + print("The newline above me should be deleted!") +- +- while True: -- print("The newline above me should be deleted!") - - --while True: -+ print("The newlines above me should be deleted!") -+while True: - - --while True: +- + while True: while False: -- print("The newlines above me should be deleted!") -- + print("The newlines above me should be deleted!") ++with open("/path/to/file.txt", mode="w") as file: -+ print("The newlines above me should be deleted!") - with open("/path/to/file.txt", mode="w") as file: -- file.write("The new line above me is about to be removed!") - - + file.write("The new line above me is about to be removed!") with open("/path/to/file.txt", mode="w") as file: -- file.write("The new lines above me is about to be removed!") +- file.write("The new line above me is about to be removed!") +-with open("/path/to/file.txt", mode="w") as file: + -+ file.write("The new lines above me is about to be removed!") - with open("/path/to/file.txt", mode="r") as read_file: -+ + file.write("The new lines above me is about to be removed!") ++with open("/path/to/file.txt", mode="r") as read_file: + +- +-with open("/path/to/file.txt", mode="r") as read_file: with open("/path/to/output_file.txt", mode="w") as write_file: + write_file.writelines(read_file.readlines()) @@ -278,17 +273,11 @@ if random.randint(0, 3) == 0: if random.uniform(0, 1) > 0.5: print("Two lines above me are about to be removed!") while True: - print("The newline above me should be deleted!") while True: - - - print("The newlines above me should be deleted!") while True: - while False: - print("The newlines above me should be deleted!") with open("/path/to/file.txt", mode="w") as file: diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__ruff_test__while_py.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__ruff_test__while_py.snap new file mode 100644 index 0000000000..03203dfbc6 --- /dev/null +++ b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__ruff_test__while_py.snap @@ -0,0 +1,70 @@ +--- +source: crates/ruff_python_formatter/src/lib.rs +expression: snapshot +--- +## Input +```py +while 34: # trailing test comment + pass # trailing last statement comment + + # trailing while body comment + +# leading else comment + +else: # trailing else comment + pass + + # trailing else body comment + + +while aVeryLongConditionThatSpillsOverToTheNextLineBecauseItIsExtremelyLongAndGoesOnAndOnAndOnAndOnAndOnAndOnAndOnAndOnAndOn: # trailing comment + pass + +else: + ... + +while ( + some_condition(unformatted, args) and anotherCondition or aThirdCondition +): # comment + print("Do something") + + +while ( + some_condition(unformatted, args) # trailing some condition + and anotherCondition or aThirdCondition # trailing third condition +): # comment + print("Do something") +``` + + + +## Output +```py +while 34: # trailing test comment + pass # trailing last statement comment + + # trailing while body comment + +# leading else comment + +else: # trailing else comment + pass + + # trailing else body comment +while ( + aVeryLongConditionThatSpillsOverToTheNextLineBecauseItIsExtremelyLongAndGoesOnAndOnAndOnAndOnAndOnAndOnAndOnAndOnAndOn +): # trailing comment + pass + +else: + ... +while some_condition(unformatted, args) and anotherCondition or aThirdCondition: # comment + print("Do something") +while ( + some_condition(unformatted, args) # trailing some condition + and anotherCondition or aThirdCondition # trailing third condition +): # comment + print("Do something") +``` + + diff --git a/crates/ruff_python_formatter/src/statement/stmt_while.rs b/crates/ruff_python_formatter/src/statement/stmt_while.rs index c11cd9cbf4..7ff55a21de 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_while.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_while.rs @@ -1,12 +1,68 @@ -use crate::{verbatim_text, FormatNodeRule, PyFormatter}; -use ruff_formatter::{write, Buffer, FormatResult}; -use rustpython_parser::ast::StmtWhile; +use crate::comments::{leading_alternate_branch_comments, trailing_comments}; +use crate::expression::maybe_parenthesize::maybe_parenthesize; +use crate::prelude::*; +use crate::FormatNodeRule; +use ruff_formatter::write; +use ruff_python_ast::node::AstNode; +use rustpython_parser::ast::{Ranged, Stmt, StmtWhile}; #[derive(Default)] pub struct FormatStmtWhile; impl FormatNodeRule for FormatStmtWhile { fn fmt_fields(&self, item: &StmtWhile, f: &mut PyFormatter) -> FormatResult<()> { - write!(f, [verbatim_text(item.range)]) + let StmtWhile { + range: _, + test, + body, + orelse, + } = item; + + let comments = f.context().comments().clone(); + let dangling_comments = comments.dangling_comments(item.as_any_node_ref()); + + let body_start = body.first().map_or(test.end(), Stmt::start); + let or_else_comments_start = + dangling_comments.partition_point(|comment| comment.slice().end() < body_start); + + let (trailing_condition_comments, or_else_comments) = + dangling_comments.split_at(or_else_comments_start); + + write!( + f, + [ + text("while"), + space(), + maybe_parenthesize(test), + text(":"), + trailing_comments(trailing_condition_comments), + block_indent(&body.format()) + ] + )?; + + if !orelse.is_empty() { + // Split between leading comments before the `else` keyword and end of line comments at the end of + // the `else:` line. + let trailing_start = + or_else_comments.partition_point(|comment| comment.position().is_own_line()); + let (leading, trailing) = or_else_comments.split_at(trailing_start); + + write!( + f, + [ + leading_alternate_branch_comments(leading, body.last()), + text("else:"), + trailing_comments(trailing), + block_indent(&orelse.format()) + ] + )?; + } + + Ok(()) + } + + fn fmt_dangling_comments(&self, _node: &StmtWhile, _f: &mut PyFormatter) -> FormatResult<()> { + // Handled in `fmt_fields` + Ok(()) } } diff --git a/crates/ruff_python_formatter/src/statement/suite.rs b/crates/ruff_python_formatter/src/statement/suite.rs index 86c7654c71..7f53ec0da3 100644 --- a/crates/ruff_python_formatter/src/statement/suite.rs +++ b/crates/ruff_python_formatter/src/statement/suite.rs @@ -28,10 +28,15 @@ impl Default for FormatSuite { impl FormatRule> for FormatSuite { fn fmt(&self, statements: &Suite, f: &mut PyFormatter) -> FormatResult<()> { - let mut joiner = f.join_nodes(match self.level { + let node_level = match self.level { SuiteLevel::TopLevel => NodeLevel::TopLevel, - SuiteLevel::Nested => NodeLevel::Statement, - }); + SuiteLevel::Nested => NodeLevel::CompoundStatement, + }; + + let saved_level = f.context().node_level(); + f.context_mut().set_node_level(node_level); + + let mut joiner = f.join_nodes(node_level); let mut iter = statements.iter(); let Some(first) = iter.next() else { @@ -67,7 +72,11 @@ impl FormatRule> for FormatSuite { is_last_function_or_class_definition = is_current_function_or_class_definition; } - joiner.finish() + let result = joiner.finish(); + + f.context_mut().set_node_level(saved_level); + + result } } diff --git a/crates/ruff_python_formatter/src/trivia.rs b/crates/ruff_python_formatter/src/trivia.rs index db99c85198..e1224883b0 100644 --- a/crates/ruff_python_formatter/src/trivia.rs +++ b/crates/ruff_python_formatter/src/trivia.rs @@ -41,7 +41,6 @@ pub(crate) fn find_first_non_trivia_character_in_range( } /// Returns the number of newlines between `offset` and the first non whitespace character in the source code. -#[allow(unused)] // TODO(micha) Remove after using for statements. pub(crate) fn lines_before(code: &str, offset: TextSize) -> u32 { let head = &code[TextRange::up_to(offset)]; let mut newlines = 0u32;