From c89eab90585e431e49c7ece3f54683ac05e4cb90 Mon Sep 17 00:00:00 2001 From: David Szotten Date: Wed, 12 Jul 2023 09:29:27 +0100 Subject: [PATCH 1/3] add TokenKind::Walrus --- ...matter__trivia__tests__tokenize_colon.snap | 18 ++++++++++ ...atter__trivia__tests__tokenize_walrus.snap | 14 ++++++++ crates/ruff_python_formatter/src/trivia.rs | 36 +++++++++++++++++++ 3 files changed, 68 insertions(+) create mode 100644 crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__trivia__tests__tokenize_colon.snap create mode 100644 crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__trivia__tests__tokenize_walrus.snap diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__trivia__tests__tokenize_colon.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__trivia__tests__tokenize_colon.snap new file mode 100644 index 0000000000000..df653c91d39c2 --- /dev/null +++ b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__trivia__tests__tokenize_colon.snap @@ -0,0 +1,18 @@ +--- +source: crates/ruff_python_formatter/src/trivia.rs +expression: test_case.tokens() +--- +[ + Token { + kind: Colon, + range: 0..1, + }, + Token { + kind: Colon, + range: 1..2, + }, + Token { + kind: Colon, + range: 2..3, + }, +] diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__trivia__tests__tokenize_walrus.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__trivia__tests__tokenize_walrus.snap new file mode 100644 index 0000000000000..69ffe6acd9cc0 --- /dev/null +++ b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__trivia__tests__tokenize_walrus.snap @@ -0,0 +1,14 @@ +--- +source: crates/ruff_python_formatter/src/trivia.rs +expression: test_case.tokens() +--- +[ + Token { + kind: Walrus, + range: 0..2, + }, + Token { + kind: Walrus, + range: 2..4, + }, +] diff --git a/crates/ruff_python_formatter/src/trivia.rs b/crates/ruff_python_formatter/src/trivia.rs index 4d8a1fa8cd20f..0197360cfc90a 100644 --- a/crates/ruff_python_formatter/src/trivia.rs +++ b/crates/ruff_python_formatter/src/trivia.rs @@ -185,6 +185,10 @@ pub(crate) enum TokenKind { /// `.`. Dot, + /// `:=`. + /// TODO: name? + Walrus, + /// `else` Else, @@ -322,6 +326,15 @@ impl<'a> SimpleTokenizer<'a> { '\\' => TokenKind::Continuation, + ':' => { + if self.cursor.first() == '=' { + self.cursor.bump(); + TokenKind::Walrus + } else { + TokenKind::Colon + } + } + c => { let kind = if is_identifier_start(c) { self.cursor.eat_while(is_identifier_continuation); @@ -457,6 +470,9 @@ impl<'a> SimpleTokenizer<'a> { self.cursor = savepoint; TokenKind::Other } + } else if c == '=' && self.cursor.last() == ':' { + self.cursor.bump_back(); + TokenKind::Walrus } else { TokenKind::from_non_trivia_char(c) }; @@ -687,6 +703,26 @@ mod tests { test_case.assert_reverse_tokenization(); } + #[test] + fn tokenize_colon() { + let source = ":::"; + + let test_case = tokenize(source); + + assert_debug_snapshot!(test_case.tokens()); + test_case.assert_reverse_tokenization(); + } + + #[test] + fn tokenize_walrus() { + let source = ":=:="; + + let test_case = tokenize(source); + + assert_debug_snapshot!(test_case.tokens()); + test_case.assert_reverse_tokenization(); + } + #[test] fn tokenize_continuation() { let source = "( \\\n )"; From 06fef9fd3bfcde99f57a485a738798c3f8877f70 Mon Sep 17 00:00:00 2001 From: David Szotten Date: Wed, 12 Jul 2023 09:30:28 +0100 Subject: [PATCH 2/3] handle comments around walrus op --- .../fixtures/ruff/expression/named_expr.py | 8 +++- .../src/comments/placement.rs | 47 +++++++++++++++++++ .../src/expression/expr_named_expr.rs | 33 ++++++++----- .../format@expression__named_expr.py.snap | 21 +++++++-- 4 files changed, 94 insertions(+), 15 deletions(-) diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/named_expr.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/named_expr.py index 9377e9704e721..8c43f6dfc25fd 100644 --- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/named_expr.py +++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/named_expr.py @@ -3,10 +3,16 @@ if ( # 1 x # 2 + # 2.5 := # 3 + # 3.5 y # 4 ): - pass + ... + +# should not add brackets +if x := y: + ... y0 = (y1 := f(x)) diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index b82d4947c4edb..11d584edfcfa3 100644 --- a/crates/ruff_python_formatter/src/comments/placement.rs +++ b/crates/ruff_python_formatter/src/comments/placement.rs @@ -40,6 +40,7 @@ pub(super) fn place_comment<'a>( handle_expr_if_comment, handle_comprehension_comment, handle_trailing_expression_starred_star_end_of_line_comment, + handle_comment_before_walrus, ]; for handler in HANDLERS { comment = match handler(comment, locator) { @@ -1232,6 +1233,52 @@ fn handle_trailing_expression_starred_star_end_of_line_comment<'a>( CommentPlacement::leading(starred.as_any_node_ref(), comment) } +/// Assign comments between the target and `:=` as trailing the target instead of leading the value +/// +/// ```python +/// if ( +/// x +/// # make trailing x (instead of leading y) +/// := +/// y +/// ) +/// ... +/// ``` +// TODO: can this also handle comments before colon in dicts +fn handle_comment_before_walrus<'a>( + comment: DecoratedComment<'a>, + locator: &Locator, +) -> CommentPlacement<'a> { + let AnyNodeRef::ExprNamedExpr(named) = comment.enclosing_node() else { + return CommentPlacement::Default(comment); + }; + + if !comment.line_position().is_own_line() { + return CommentPlacement::Default(comment); + } + + let walrus = find_only_token_in_range( + TextRange::new(named.target.range().end(), named.value.range().start()), + locator, + TokenKind::Walrus, + ); + + // Comments between the target and the `:=` + // ```python + // [ + // a + // # attach as trailing on the target `a` (instead of leading on the value `b`) + // := + // in b + // ] + // ``` + if comment.slice().end() < walrus.start() { + return CommentPlacement::trailing((&*named.target).into(), comment); + } + + CommentPlacement::Default(comment) +} + /// Looks for a token in the range that contains no other tokens except for parentheses outside /// the expression ranges fn find_only_token_in_range(range: TextRange, locator: &Locator, token_kind: TokenKind) -> Token { diff --git a/crates/ruff_python_formatter/src/expression/expr_named_expr.rs b/crates/ruff_python_formatter/src/expression/expr_named_expr.rs index e1ba5c1789e9f..a3a391bb7a027 100644 --- a/crates/ruff_python_formatter/src/expression/expr_named_expr.rs +++ b/crates/ruff_python_formatter/src/expression/expr_named_expr.rs @@ -3,7 +3,7 @@ use crate::expression::parentheses::{ default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize, }; use crate::{AsFormat, FormatNodeRule, PyFormatter}; -use ruff_formatter::prelude::{space, text}; +use ruff_formatter::prelude::{soft_line_break, space, text}; use ruff_formatter::{write, Buffer, FormatResult}; use rustpython_parser::ast::ExprNamedExpr; @@ -17,16 +17,27 @@ impl FormatNodeRule for FormatExprNamedExpr { value, range: _, } = item; - write!( - f, - [ - target.format(), - space(), - text(":="), - space(), - value.format(), - ] - ) + + write!(f, [target.format()])?; + + let comments = f.context().comments().clone(); + let trailing_target_comments = comments.trailing_comments(target.as_ref()); + let leading_value_comments = comments.leading_comments(value.as_ref()); + + if trailing_target_comments.is_empty() { + write!(f, [space()])?; + } else { + write!(f, [soft_line_break()])?; + } + + write!(f, [text(":=")])?; + + if leading_value_comments.is_empty() { + write!(f, [space()])?; + } else { + write!(f, [soft_line_break()])?; + } + write!(f, [value.format()]) } } diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap index 9cf0484b4bbaa..2503ff62374b5 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap +++ b/crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap @@ -9,10 +9,16 @@ y = 1 if ( # 1 x # 2 + # 2.5 := # 3 + # 3.5 y # 4 ): - pass + ... + +# should not add brackets +if x := y: + ... y0 = (y1 := f(x)) @@ -25,9 +31,18 @@ y = 1 if ( # 1 - x := y # 2 # 3 # 4 + x # 2 + # 2.5 + # 3 + := + # 3.5 + y # 4 ): - pass + ... + +# should not add brackets +if (x := y): + ... y0 = (y1 := f(x)) From 982b712a7d8c81c0b573c6e2cf70f3dddf9f469b Mon Sep 17 00:00:00 2001 From: David Szotten Date: Wed, 12 Jul 2023 09:45:37 +0100 Subject: [PATCH 3/3] also handle end of line comments trailing a lone walrus --- .../src/comments/placement.rs | 36 ++++++++++++++----- .../src/expression/expr_named_expr.rs | 13 ++++++- .../format@expression__named_expr.py.snap | 3 +- 3 files changed, 40 insertions(+), 12 deletions(-) diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index 11d584edfcfa3..2faaec8adca99 100644 --- a/crates/ruff_python_formatter/src/comments/placement.rs +++ b/crates/ruff_python_formatter/src/comments/placement.rs @@ -40,7 +40,7 @@ pub(super) fn place_comment<'a>( handle_expr_if_comment, handle_comprehension_comment, handle_trailing_expression_starred_star_end_of_line_comment, - handle_comment_before_walrus, + handle_walrus_comments, ]; for handler in HANDLERS { comment = match handler(comment, locator) { @@ -1233,19 +1233,19 @@ fn handle_trailing_expression_starred_star_end_of_line_comment<'a>( CommentPlacement::leading(starred.as_any_node_ref(), comment) } -/// Assign comments between the target and `:=` as trailing the target instead of leading the value +/// Handle comments around `:=` /// /// ```python /// if ( /// x /// # make trailing x (instead of leading y) -/// := +/// := # make dangling on the NamedExpr node /// y /// ) /// ... /// ``` // TODO: can this also handle comments before colon in dicts -fn handle_comment_before_walrus<'a>( +fn handle_walrus_comments<'a>( comment: DecoratedComment<'a>, locator: &Locator, ) -> CommentPlacement<'a> { @@ -1253,16 +1253,14 @@ fn handle_comment_before_walrus<'a>( return CommentPlacement::Default(comment); }; - if !comment.line_position().is_own_line() { - return CommentPlacement::Default(comment); - } - let walrus = find_only_token_in_range( TextRange::new(named.target.range().end(), named.value.range().start()), locator, TokenKind::Walrus, ); + let is_own_line = comment.line_position().is_own_line(); + // Comments between the target and the `:=` // ```python // [ @@ -1273,7 +1271,27 @@ fn handle_comment_before_walrus<'a>( // ] // ``` if comment.slice().end() < walrus.start() { - return CommentPlacement::trailing((&*named.target).into(), comment); + return if is_own_line { + CommentPlacement::trailing((&*named.target).into(), comment) + } else { + CommentPlacement::Default(comment) + }; + } + + // Comments after the `:=` + // ```python + // [ + // a + // := # attach as dangling on the NamedExpr + // in b + // ] + // ``` + if comment.slice().end() < named.value.range().start() { + return if is_own_line { + CommentPlacement::Default(comment) + } else { + return CommentPlacement::dangling(comment.enclosing_node(), comment); + }; } CommentPlacement::Default(comment) diff --git a/crates/ruff_python_formatter/src/expression/expr_named_expr.rs b/crates/ruff_python_formatter/src/expression/expr_named_expr.rs index a3a391bb7a027..7e56850ffb479 100644 --- a/crates/ruff_python_formatter/src/expression/expr_named_expr.rs +++ b/crates/ruff_python_formatter/src/expression/expr_named_expr.rs @@ -1,3 +1,4 @@ +use crate::comments::dangling_comments; use crate::context::PyFormatContext; use crate::expression::parentheses::{ default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize, @@ -23,6 +24,7 @@ impl FormatNodeRule for FormatExprNamedExpr { let comments = f.context().comments().clone(); let trailing_target_comments = comments.trailing_comments(target.as_ref()); let leading_value_comments = comments.leading_comments(value.as_ref()); + let dangling_item_comments = comments.dangling_comments(item); if trailing_target_comments.is_empty() { write!(f, [space()])?; @@ -30,7 +32,7 @@ impl FormatNodeRule for FormatExprNamedExpr { write!(f, [soft_line_break()])?; } - write!(f, [text(":=")])?; + write!(f, [text(":="), dangling_comments(dangling_item_comments)])?; if leading_value_comments.is_empty() { write!(f, [space()])?; @@ -39,6 +41,15 @@ impl FormatNodeRule for FormatExprNamedExpr { } write!(f, [value.format()]) } + + fn fmt_dangling_comments( + &self, + _node: &ExprNamedExpr, + _f: &mut PyFormatter, + ) -> FormatResult<()> { + // Handled in `fmt_fields` + Ok(()) + } } impl NeedsParentheses for ExprNamedExpr { diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap index 2503ff62374b5..c327062c58f0c 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap +++ b/crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap @@ -33,8 +33,7 @@ if ( # 1 x # 2 # 2.5 - # 3 - := + := # 3 # 3.5 y # 4 ):