From 8f9888498604e3dfb200fce17ba4919d771dfabd Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Fri, 14 Jul 2023 09:56:38 +0200 Subject: [PATCH] Parenthesize with statements --- .../test/fixtures/ruff/statement/with.py | 62 +++++++- crates/ruff_python_formatter/src/builders.rs | 15 +- .../src/comments/placement.rs | 45 ++++++ .../src/expression/mod.rs | 29 ++-- .../src/expression/parentheses.rs | 5 + crates/ruff_python_formatter/src/lib.rs | 33 +++-- .../src/other/with_item.rs | 38 +++-- .../src/statement/stmt_with.rs | 77 ++++++++-- crates/ruff_python_formatter/src/trivia.rs | 14 +- .../snapshots/format@statement__with.py.snap | 140 ++++++++++++++++-- 10 files changed, 392 insertions(+), 66 deletions(-) diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with.py index 4a664a1d70d784..f959e005a7fa02 100644 --- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with.py +++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with.py @@ -40,9 +40,6 @@ with (a): # should remove brackets ... -# TODO: black doesn't wrap this, but maybe we want to anyway? -# if we do want to wrap, do we prefer to wrap the entire WithItem or to let the -# WithItem allow the `aa + bb` content expression to be wrapped with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as c: ... @@ -52,3 +49,62 @@ pass with (a, *b): pass + +with ( + # leading comment + a) as b: ... + +with ( + # leading comment + a as b +): ... + + +with (a # trailing same line comment + # trailing own line comment + ) as b: ... + +with ( + a # trailing same line comment + # trailing own line comment + as b +): ... + + +with (a # trailing same line comment + # trailing own line comment +) as b: ... + +with ( + (a + # trailing own line comment + ) + as # trailing as same line comment + b # trailing b same line comment +): ... + +with ( + [ + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "bbbbbbbbbb", + "cccccccccccccccccccccccccccccccccccccccccc", + dddddddddddddddddddddddddddddddd, + ] as example1, + aaaaaaaaaaaaaaaaaaaaaaaaaa + + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb + + cccccccccccccccccccccccccccc + + ddddddddddddddddd as example2, + CtxManager2() as example2, + CtxManager2() as example2, + CtxManager2() as example2, +): + ... + + +with [ + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "bbbbbbbbbb", + "cccccccccccccccccccccccccccccccccccccccccc", + dddddddddddddddddddddddddddddddd, +] as example1, aaaaaaaaaaaaaaaaaaaaaaaaaa * bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb * cccccccccccccccccccccccccccc + ddddddddddddddddd as example2, CtxManager222222222222222() as example2: + ... diff --git a/crates/ruff_python_formatter/src/builders.rs b/crates/ruff_python_formatter/src/builders.rs index 5ff156917041ee..9f24a49ca2a7d1 100644 --- a/crates/ruff_python_formatter/src/builders.rs +++ b/crates/ruff_python_formatter/src/builders.rs @@ -227,10 +227,23 @@ impl<'fmt, 'ast, 'buf> JoinCommaSeparatedBuilder<'fmt, 'ast, 'buf> { ) -> &mut Self where T: Ranged, + { + self.entry_with_line_separator(node, content, soft_line_break_or_space()) + } + + pub(crate) fn entry_with_line_separator( + &mut self, + node: &N, + content: &dyn Format>, + separator: Separator, + ) -> &mut Self + where + N: Ranged, + Separator: Format>, { self.result = self.result.and_then(|_| { if self.end_of_last_entry.is_some() { - write!(self.fmt, [text(","), soft_line_break_or_space()])?; + write!(self.fmt, [text(","), separator])?; } self.end_of_last_entry = Some(node.end()); diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index b82d4947c4edb8..6effa9e3f5efb3 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_with_item_comment, ]; for handler in HANDLERS { comment = match handler(comment, locator) { @@ -1232,6 +1233,50 @@ fn handle_trailing_expression_starred_star_end_of_line_comment<'a>( CommentPlacement::leading(starred.as_any_node_ref(), comment) } +/// Handles trailing own line comments before the `as` keyword of a with item and +/// end of line comments that are on the same line as the `as` keyword: +/// +/// ```python +/// with ( +/// a +/// # trailing a own line comment +/// as # trailing as same line comment +/// b +// ): ... +/// ``` +fn handle_with_item_comment<'a>( + comment: DecoratedComment<'a>, + locator: &Locator, +) -> CommentPlacement<'a> { + if !comment.enclosing_node().is_with_item() { + return CommentPlacement::Default(comment); + } + + // Needs to be a with item with an `as` expression. + let (Some(context_expr), Some(optional_vars)) = + (comment.preceding_node(), comment.following_node()) + else { + return CommentPlacement::Default(comment); + }; + + let as_token = find_only_token_in_range( + TextRange::new(context_expr.end(), optional_vars.start()), + locator, + TokenKind::As, + ); + + // If before the `as` keyword, then it must be a trailing comment of the context expression. + if comment.end() < as_token.start() { + CommentPlacement::trailing(context_expr, comment) + } + // Trailing end of line comment coming after the `as` keyword`. + else if comment.line_position().is_end_of_line() { + CommentPlacement::dangling(comment.enclosing_node(), comment) + } else { + 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/mod.rs b/crates/ruff_python_formatter/src/expression/mod.rs index 617fd61c41629e..0ac1681b1e45c6 100644 --- a/crates/ruff_python_formatter/src/expression/mod.rs +++ b/crates/ruff_python_formatter/src/expression/mod.rs @@ -159,22 +159,29 @@ impl Format> for MaybeParenthesizeExpression<'_> { parenthesize, } = self; - let parenthesize = match parenthesize { + let preserve_parentheses = match parenthesize { Parenthesize::Optional => { is_expression_parenthesized(AnyNodeRef::from(*expression), f.context().source()) } - Parenthesize::IfBreaks => false, + Parenthesize::IfBreaks | Parenthesize::IfRequired => false, }; - let parentheses = - if parenthesize || f.context().comments().has_leading_comments(*expression) { - OptionalParentheses::Always - } else { - expression.needs_parentheses(*parent, f.context()) - }; + let comments = f.context().comments(); + let needs_parentheses = if preserve_parentheses + || comments.has_leading_comments(*expression) + || comments.has_trailing_own_line_comments(*expression) + { + OptionalParentheses::Always + } else if *parenthesize == Parenthesize::IfRequired + && f.context().node_level().is_parenthesized() + { + OptionalParentheses::Never + } else { + expression.needs_parentheses(*parent, f.context()) + }; - match parentheses { - OptionalParentheses::Multiline => { + match needs_parentheses { + OptionalParentheses::Multiline if *parenthesize != Parenthesize::IfRequired => { if can_omit_optional_parentheses(expression, f.context()) { optional_parentheses(&expression.format().with_options(Parentheses::Never)) .fmt(f) @@ -186,7 +193,7 @@ impl Format> for MaybeParenthesizeExpression<'_> { OptionalParentheses::Always => { expression.format().with_options(Parentheses::Always).fmt(f) } - OptionalParentheses::Never => { + OptionalParentheses::Never | OptionalParentheses::Multiline => { expression.format().with_options(Parentheses::Never).fmt(f) } } diff --git a/crates/ruff_python_formatter/src/expression/parentheses.rs b/crates/ruff_python_formatter/src/expression/parentheses.rs index f6ea6a22db1651..652a05ba6a155f 100644 --- a/crates/ruff_python_formatter/src/expression/parentheses.rs +++ b/crates/ruff_python_formatter/src/expression/parentheses.rs @@ -36,6 +36,11 @@ pub(crate) enum Parenthesize { /// Parenthesizes the expression only if it doesn't fit on a line. IfBreaks, + + /// Only adds parentheses if absolutely necessary: + /// * The expression is not enclosed by another parenthesized expression and it expands over multiple lines + /// * The expression has leading or trailing comments. Adding parentheses is desired to prevent the comments from wandering. + IfRequired, } /// Whether it is necessary to add parentheses around an expression. diff --git a/crates/ruff_python_formatter/src/lib.rs b/crates/ruff_python_formatter/src/lib.rs index 44ab2e8de4e7af..7055bab2faf569 100644 --- a/crates/ruff_python_formatter/src/lib.rs +++ b/crates/ruff_python_formatter/src/lib.rs @@ -280,15 +280,30 @@ if True: #[test] fn quick_test() { let src = r#" -if a * [ - bbbbbbbbbbbbbbbbbbbbbb, - cccccccccccccccccccccccccccccdddddddddddddddddddddddddd, -] + a * e * [ - ffff, - gggg, - hhhhhhhhhhhhhh, -] * c: - pass +with ( + [ + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "bbbbbbbbbb", + "cccccccccccccccccccccccccccccccccccccccccc", + dddddddddddddddddddddddddddddddd, + ] as example1, + aaaaaaaaaaaaaaaaaaaaaaaaaa + + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb + + cccccccccccccccccccccccccccc + + ddddddddddddddddd as example2, + CtxManager2() as example2, + CtxManager2() as example2, + CtxManager2() as example2, +): + ... + +with [ + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "bbbbbbbbbb", + "cccccccccccccccccccccccccccccccccccccccccc", + dddddddddddddddddddddddddddddddd, +] as example1, aaaaaaaaaaaaaaaaaaaaaaaaaa * bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb * cccccccccccccccccccccccccccc + ddddddddddddddddd as example2, CtxManager222222222222222() as example2: + ... "#; // Tokenize once diff --git a/crates/ruff_python_formatter/src/other/with_item.rs b/crates/ruff_python_formatter/src/other/with_item.rs index 1004639f15f289..6a4034d6331fb6 100644 --- a/crates/ruff_python_formatter/src/other/with_item.rs +++ b/crates/ruff_python_formatter/src/other/with_item.rs @@ -1,9 +1,12 @@ +use rustpython_parser::ast::WithItem; + +use ruff_formatter::{write, Buffer, FormatResult}; + +use crate::comments::trailing_comments; use crate::expression::maybe_parenthesize_expression; use crate::expression::parentheses::Parenthesize; use crate::prelude::*; use crate::{FormatNodeRule, PyFormatter}; -use ruff_formatter::{write, Buffer, FormatResult}; -use rustpython_parser::ast::WithItem; #[derive(Default)] pub struct FormatWithItem; @@ -16,20 +19,27 @@ impl FormatNodeRule for FormatWithItem { optional_vars, } = item; - let inner = format_with(|f| { + let comments = f.context().comments().clone(); + let trailing_as_comments = comments.dangling_comments(item); + + maybe_parenthesize_expression(context_expr, item, Parenthesize::IfRequired).fmt(f)?; + + if let Some(optional_vars) = optional_vars { write!( f, - [maybe_parenthesize_expression( - context_expr, - item, - Parenthesize::IfBreaks - )] + [ + space(), + text("as"), + trailing_comments(trailing_as_comments), + space(), + optional_vars.format(), + ] )?; - if let Some(optional_vars) = optional_vars { - write!(f, [space(), text("as"), space(), optional_vars.format()])?; - } - Ok(()) - }); - write!(f, [group(&inner)]) + } + Ok(()) + } + + fn fmt_dangling_comments(&self, _node: &WithItem, _f: &mut PyFormatter) -> FormatResult<()> { + Ok(()) } } diff --git a/crates/ruff_python_formatter/src/statement/stmt_with.rs b/crates/ruff_python_formatter/src/statement/stmt_with.rs index 8474d702401b5e..56eca66b17236d 100644 --- a/crates/ruff_python_formatter/src/statement/stmt_with.rs +++ b/crates/ruff_python_formatter/src/statement/stmt_with.rs @@ -1,11 +1,15 @@ -use ruff_formatter::{write, Buffer, FormatResult}; -use ruff_python_ast::node::AnyNodeRef; use ruff_text_size::TextRange; use rustpython_parser::ast::{Ranged, StmtAsyncWith, StmtWith, Suite, WithItem}; -use crate::builders::parenthesize_if_expands; +use ruff_formatter::{format_args, write, FormatError}; +use ruff_python_ast::node::AnyNodeRef; + use crate::comments::trailing_comments; +use crate::expression::parentheses::{ + in_parentheses_only_soft_line_break_or_space, optional_parentheses, +}; use crate::prelude::*; +use crate::trivia::{SimpleTokenizer, TokenKind}; use crate::FormatNodeRule; pub(super) enum AnyStatementWith<'a> { @@ -68,22 +72,39 @@ impl Format> for AnyStatementWith<'_> { let comments = f.context().comments().clone(); let dangling_comments = comments.dangling_comments(self); - let joined_items = format_with(|f| { - f.join_comma_separated(self.body().first().unwrap().start()) - .nodes(self.items().iter()) - .finish() - }); - - if self.is_async() { - write!(f, [text("async"), space()])?; + write!( + f, + [ + self.is_async() + .then_some(format_args![text("async"), space()]), + text("with"), + space() + ] + )?; + + if are_with_items_parenthesized(self, f.context())? { + optional_parentheses(&format_with(|f| { + let mut joiner = f.join_comma_separated(self.body().first().unwrap().start()); + + for item in self.items() { + joiner.entry_with_line_separator( + item, + &item.format(), + in_parentheses_only_soft_line_break_or_space(), + ); + } + joiner.finish() + })) + .fmt(f)?; + } else { + f.join_with(format_args![text(","), space()]) + .entries(self.items().iter().formatted()) + .finish()?; } write!( f, [ - text("with"), - space(), - group(&parenthesize_if_expands(&joined_items)), text(":"), trailing_comments(dangling_comments), block_indent(&self.body().format()) @@ -92,6 +113,34 @@ impl Format> for AnyStatementWith<'_> { } } +fn are_with_items_parenthesized( + with: &AnyStatementWith, + context: &PyFormatContext, +) -> FormatResult { + let first_with_item = with.items().first().ok_or(FormatError::SyntaxError)?; + let before_first_with_item = TextRange::new(with.start(), first_with_item.start()); + + let mut tokenizer = SimpleTokenizer::new(context.source(), before_first_with_item) + .skip_trivia() + .skip_while(|t| t.kind() == TokenKind::Async); + + let with_keyword = tokenizer.next().ok_or(FormatError::SyntaxError)?; + + debug_assert_eq!( + with_keyword.kind(), + TokenKind::With, + "Expected with keyword but at {with_keyword:?}" + ); + + match tokenizer.next() { + Some(left_paren) => { + debug_assert_eq!(left_paren.kind(), TokenKind::LParen); + Ok(true) + } + None => Ok(false), + } +} + #[derive(Default)] pub struct FormatStmtWith; diff --git a/crates/ruff_python_formatter/src/trivia.rs b/crates/ruff_python_formatter/src/trivia.rs index b6432e53364e65..b1396fad72d587 100644 --- a/crates/ruff_python_formatter/src/trivia.rs +++ b/crates/ruff_python_formatter/src/trivia.rs @@ -193,9 +193,18 @@ pub(crate) enum TokenKind { /// `in` In, + /// `as` + As, + /// `match` Match, + /// `with` + With, + + /// `async` + Async, + /// Any other non trivia token. Other, @@ -272,10 +281,13 @@ impl<'a> SimpleTokenizer<'a> { fn to_keyword_or_other(&self, range: TextRange) -> TokenKind { let source = &self.source[range]; match source { - "if" => TokenKind::If, + "as" => TokenKind::As, + "async" => TokenKind::Async, "else" => TokenKind::Else, + "if" => TokenKind::If, "in" => TokenKind::In, "match" => TokenKind::Match, // Match is a soft keyword that depends on the context but we can always lex it as a keyword and leave it to the caller (parser) to decide if it should be handled as an identifier or keyword. + "with" => TokenKind::With, // ..., _ => TokenKind::Other, // Potentially an identifier, but only if it isn't a string prefix. We can ignore this for now https://docs.python.org/3/reference/lexical_analysis.html#string-and-bytes-literals } diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__with.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__with.py.snap index 2d133fe87be6be..38b70eb2adbc49 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@statement__with.py.snap +++ b/crates/ruff_python_formatter/tests/snapshots/format@statement__with.py.snap @@ -46,9 +46,6 @@ with (a,): # magic trailing comma with (a): # should remove brackets ... -# TODO: black doesn't wrap this, but maybe we want to anyway? -# if we do want to wrap, do we prefer to wrap the entire WithItem or to let the -# WithItem allow the `aa + bb` content expression to be wrapped with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as c: ... @@ -58,14 +55,70 @@ with (name_2 for name_0 in name_4): pass with (a, *b): pass + +with ( + # leading comment + a) as b: ... + +with ( + # leading comment + a as b +): ... + + +with (a # trailing same line comment + # trailing own line comment + ) as b: ... + +with ( + a # trailing same line comment + # trailing own line comment + as b +): ... + + +with (a # trailing same line comment + # trailing own line comment +) as b: ... + +with ( + (a + # trailing own line comment + ) + as # trailing as same line comment + b # trailing b same line comment +): ... + +with ( + [ + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "bbbbbbbbbb", + "cccccccccccccccccccccccccccccccccccccccccc", + dddddddddddddddddddddddddddddddd, + ] as example1, + aaaaaaaaaaaaaaaaaaaaaaaaaa + + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb + + cccccccccccccccccccccccccccc + + ddddddddddddddddd as example2, + CtxManager2() as example2, + CtxManager2() as example2, + CtxManager2() as example2, +): + ... + + +with [ + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "bbbbbbbbbb", + "cccccccccccccccccccccccccccccccccccccccccc", + dddddddddddddddddddddddddddddddd, +] as example1, aaaaaaaaaaaaaaaaaaaaaaaaaa * bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb * cccccccccccccccccccccccccccc + ddddddddddddddddd as example2, CtxManager222222222222222() as example2: + ... ``` ## Output ```py -with ( - aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa, - aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa, -): +with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa, aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa: ... # trailing @@ -105,12 +158,7 @@ with ( with a: # should remove brackets ... -# TODO: black doesn't wrap this, but maybe we want to anyway? -# if we do want to wrap, do we prefer to wrap the entire WithItem or to let the -# WithItem allow the `aa + bb` content expression to be wrapped -with ( - aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as c -): +with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as c: ... @@ -119,6 +167,72 @@ with (NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in pass with (a, *b): pass + +with ( + # leading comment + a +) as b: + ... + +with ( + # leading comment + a as b +): + ... + + +with ( + a # trailing same line comment + # trailing own line comment +) as b: + ... + +with ( + a # trailing same line comment + # trailing own line comment +) as b: + ... + + +with ( + a # trailing same line comment + # trailing own line comment +) as b: + ... + +with ( + ( + a + # trailing own line comment + ) as b # trailing as same line comment # trailing b same line comment +): + ... + +with ( + [ + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "bbbbbbbbbb", + "cccccccccccccccccccccccccccccccccccccccccc", + dddddddddddddddddddddddddddddddd, + ] as example1, + aaaaaaaaaaaaaaaaaaaaaaaaaa + + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb + + cccccccccccccccccccccccccccc + + ddddddddddddddddd as example2, + CtxManager2() as example2, + CtxManager2() as example2, + CtxManager2() as example2, +): + ... + + +with [ + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "bbbbbbbbbb", + "cccccccccccccccccccccccccccccccccccccccccc", + dddddddddddddddddddddddddddddddd, +] as example1, aaaaaaaaaaaaaaaaaaaaaaaaaa * bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb * cccccccccccccccccccccccccccc + ddddddddddddddddd as example2, CtxManager222222222222222() as example2: + ... ```