Cover Black's is_aritmetic_like formatting (astral-sh#5738)

evanrittenhouse · Jul 19, 2023 · 39c87e6 · 39c87e6
1 parent f036854
commit 39c87e6
Show file tree

Hide file tree

Showing 14 changed files with 159 additions and 96 deletions.
diff --git a/crates/ruff_python_formatter/src/context.rs b/crates/ruff_python_formatter/src/context.rs
@@ -84,3 +84,13 @@ pub(crate) enum NodeLevel {
     /// Formatting nodes that are enclosed by a parenthesized (any `[]`, `{}` or `()`) expression.
     ParenthesizedExpression,
 }
+
+impl NodeLevel {
+    /// Returns `true` if the expression is in a parenthesized context.
+    pub(crate) const fn is_parenthesized(self) -> bool {
+        matches!(
+            self,
+            NodeLevel::Expression(Some(_)) | NodeLevel::ParenthesizedExpression
+        )
+    }
+}
diff --git a/crates/ruff_python_formatter/src/expression/expr_await.rs b/crates/ruff_python_formatter/src/expression/expr_await.rs
@@ -1,18 +1,29 @@
-use crate::context::PyFormatContext;
-use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
-use crate::{AsFormat, FormatNodeRule, PyFormatter};
-use ruff_formatter::prelude::{space, text};
-use ruff_formatter::{write, Buffer, FormatResult};
-use ruff_python_ast::node::AnyNodeRef;
 use rustpython_parser::ast::ExprAwait;
 
+use ruff_formatter::write;
+use ruff_python_ast::node::AnyNodeRef;
+
+use crate::expression::maybe_parenthesize_expression;
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parenthesize};
+use crate::prelude::*;
+use crate::FormatNodeRule;
+
 #[derive(Default)]
 pub struct FormatExprAwait;
 
 impl FormatNodeRule<ExprAwait> for FormatExprAwait {
     fn fmt_fields(&self, item: &ExprAwait, f: &mut PyFormatter) -> FormatResult<()> {
         let ExprAwait { range: _, value } = item;
-        write!(f, [text("await"), space(), value.format()])
+
+        let format_value = format_with(|f: &mut PyFormatter| {
+            if f.context().node_level().is_parenthesized() {
+                value.format().fmt(f)
+            } else {
+                maybe_parenthesize_expression(value, item, Parenthesize::Optional).fmt(f)
+            }
+        });
+
+        write!(f, [text("await"), space(), format_value])
     }
 }
 

diff --git a/crates/ruff_python_formatter/src/expression/expr_bin_op.rs b/crates/ruff_python_formatter/src/expression/expr_bin_op.rs
@@ -1,6 +1,8 @@
 use crate::comments::{trailing_comments, trailing_node_comments};
 use crate::expression::parentheses::{
-    in_parentheses_only_group, is_expression_parenthesized, NeedsParentheses, OptionalParentheses,
+    in_parentheses_only_group, in_parentheses_only_soft_line_break,
+    in_parentheses_only_soft_line_break_or_space, is_expression_parenthesized, NeedsParentheses,
+    OptionalParentheses,
 };
 use crate::expression::Parentheses;
 use crate::prelude::*;
@@ -64,9 +66,9 @@ impl FormatNodeRule<ExprBinOp> for FormatExprBinOp {
                 let needs_space = !is_simple_power_expression(current);
 
                 let before_operator_space = if needs_space {
-                    soft_line_break_or_space()
+                    in_parentheses_only_soft_line_break_or_space()
                 } else {
-                    soft_line_break()
+                    in_parentheses_only_soft_line_break()
                 };
 
                 write!(

diff --git a/crates/ruff_python_formatter/src/expression/expr_bool_op.rs b/crates/ruff_python_formatter/src/expression/expr_bool_op.rs
@@ -1,6 +1,7 @@
 use crate::comments::leading_comments;
 use crate::expression::parentheses::{
-    in_parentheses_only_group, NeedsParentheses, OptionalParentheses, Parentheses,
+    in_parentheses_only_group, in_parentheses_only_soft_line_break_or_space, NeedsParentheses,
+    OptionalParentheses, Parentheses,
 };
 use crate::prelude::*;
 use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions};
@@ -42,7 +43,7 @@ impl FormatNodeRule<ExprBoolOp> for FormatExprBoolOp {
                 let leading_value_comments = comments.leading_comments(value);
                 // Format the expressions leading comments **before** the operator
                 if leading_value_comments.is_empty() {
-                    write!(f, [soft_line_break_or_space()])?;
+                    write!(f, [in_parentheses_only_soft_line_break_or_space()])?;
                 } else {
                     write!(
                         f,

diff --git a/crates/ruff_python_formatter/src/expression/expr_compare.rs b/crates/ruff_python_formatter/src/expression/expr_compare.rs
@@ -1,6 +1,7 @@
 use crate::comments::leading_comments;
 use crate::expression::parentheses::{
-    in_parentheses_only_group, NeedsParentheses, OptionalParentheses, Parentheses,
+    in_parentheses_only_group, in_parentheses_only_soft_line_break_or_space, NeedsParentheses,
+    OptionalParentheses, Parentheses,
 };
 use crate::prelude::*;
 use crate::FormatNodeRule;
@@ -41,7 +42,7 @@ impl FormatNodeRule<ExprCompare> for FormatExprCompare {
             for (operator, comparator) in ops.iter().zip(comparators) {
                 let leading_comparator_comments = comments.leading_comments(comparator);
                 if leading_comparator_comments.is_empty() {
-                    write!(f, [soft_line_break_or_space()])?;
+                    write!(f, [in_parentheses_only_soft_line_break_or_space()])?;
                 } else {
                     // Format the expressions leading comments **before** the operator
                     write!(

diff --git a/crates/ruff_python_formatter/src/expression/expr_if_exp.rs b/crates/ruff_python_formatter/src/expression/expr_if_exp.rs
@@ -1,6 +1,7 @@
 use crate::comments::leading_comments;
 use crate::expression::parentheses::{
-    in_parentheses_only_group, NeedsParentheses, OptionalParentheses,
+    in_parentheses_only_group, in_parentheses_only_soft_line_break_or_space, NeedsParentheses,
+    OptionalParentheses,
 };
 use crate::prelude::*;
 use crate::FormatNodeRule;
@@ -28,12 +29,12 @@ impl FormatNodeRule<ExprIfExp> for FormatExprIfExp {
             f,
             [in_parentheses_only_group(&format_args![
                 body.format(),
-                soft_line_break_or_space(),
+                in_parentheses_only_soft_line_break_or_space(),
                 leading_comments(comments.leading_comments(test.as_ref())),
                 text("if"),
                 space(),
                 test.format(),
-                soft_line_break_or_space(),
+                in_parentheses_only_soft_line_break_or_space(),
                 leading_comments(comments.leading_comments(orelse.as_ref())),
                 text("else"),
                 space(),

diff --git a/crates/ruff_python_formatter/src/expression/expr_subscript.rs b/crates/ruff_python_formatter/src/expression/expr_subscript.rs
@@ -7,9 +7,7 @@ use crate::comments::trailing_comments;
 use crate::context::NodeLevel;
 use crate::context::PyFormatContext;
 use crate::expression::expr_tuple::TupleParentheses;
-use crate::expression::parentheses::{
-    in_parentheses_only_group, NeedsParentheses, OptionalParentheses,
-};
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::prelude::*;
 use crate::FormatNodeRule;
 
@@ -64,7 +62,7 @@ impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
 
         write!(
             f,
-            [in_parentheses_only_group(&format_args![
+            [group(&format_args![
                 text("["),
                 trailing_comments(dangling_comments),
                 soft_block_indent(&format_slice),

diff --git a/crates/ruff_python_formatter/src/expression/parentheses.rs b/crates/ruff_python_formatter/src/expression/parentheses.rs
@@ -2,7 +2,7 @@ use crate::context::NodeLevel;
 use crate::prelude::*;
 use crate::trivia::{first_non_trivia_token, first_non_trivia_token_rev, Token, TokenKind};
 use ruff_formatter::prelude::tag::Condition;
-use ruff_formatter::{format_args, Argument, Arguments};
+use ruff_formatter::{format_args, write, Argument, Arguments};
 use ruff_python_ast::node::AnyNodeRef;
 use rustpython_parser::ast::Ranged;
 
@@ -178,6 +178,55 @@ impl<'ast> Format<PyFormatContext<'ast>> for FormatOptionalParentheses<'_, 'ast>
     }
 }
 
+/// Creates a [`soft_line_break`] if the expression is enclosed by (optional) parentheses (`()`, `[]`, or `{}`).
+/// Prints nothing if the expression is not parenthesized.
+pub(crate) const fn in_parentheses_only_soft_line_break() -> InParenthesesOnlyLineBreak {
+    InParenthesesOnlyLineBreak::SoftLineBreak
+}
+
+/// Creates a [`soft_line_break_or_space`] if the expression is enclosed by (optional) parentheses (`()`, `[]`, or `{}`).
+/// Prints a [`space`] if the expression is not parenthesized.
+pub(crate) const fn in_parentheses_only_soft_line_break_or_space() -> InParenthesesOnlyLineBreak {
+    InParenthesesOnlyLineBreak::SoftLineBreakOrSpace
+}
+
+pub(crate) enum InParenthesesOnlyLineBreak {
+    SoftLineBreak,
+    SoftLineBreakOrSpace,
+}
+
+impl<'ast> Format<PyFormatContext<'ast>> for InParenthesesOnlyLineBreak {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'ast>>) -> FormatResult<()> {
+        match f.context().node_level() {
+            NodeLevel::TopLevel | NodeLevel::CompoundStatement | NodeLevel::Expression(None) => {
+                match self {
+                    InParenthesesOnlyLineBreak::SoftLineBreak => Ok(()),
+                    InParenthesesOnlyLineBreak::SoftLineBreakOrSpace => space().fmt(f),
+                }
+            }
+            NodeLevel::Expression(Some(parentheses_id)) => match self {
+                InParenthesesOnlyLineBreak::SoftLineBreak => if_group_breaks(&soft_line_break())
+                    .with_group_id(Some(parentheses_id))
+                    .fmt(f),
+                InParenthesesOnlyLineBreak::SoftLineBreakOrSpace => write!(
+                    f,
+                    [
+                        if_group_breaks(&soft_line_break_or_space())
+                            .with_group_id(Some(parentheses_id)),
+                        if_group_fits_on_line(&space()).with_group_id(Some(parentheses_id))
+                    ]
+                ),
+            },
+            NodeLevel::ParenthesizedExpression => {
+                f.write_element(FormatElement::Line(match self {
+                    InParenthesesOnlyLineBreak::SoftLineBreak => LineMode::Soft,
+                    InParenthesesOnlyLineBreak::SoftLineBreakOrSpace => LineMode::SoftOrSpace,
+                }))
+            }
+        }
+    }
+}
+
 /// Makes `content` a group, but only if the outer expression is parenthesized (a list, parenthesized expression, dict, ...)
 /// or if the expression gets parenthesized because it expands over multiple lines.
 pub(crate) fn in_parentheses_only_group<'content, 'ast, Content>(
@@ -197,17 +246,23 @@ pub(crate) struct FormatInParenthesesOnlyGroup<'content, 'ast> {
 
 impl<'ast> Format<PyFormatContext<'ast>> for FormatInParenthesesOnlyGroup<'_, 'ast> {
     fn fmt(&self, f: &mut Formatter<PyFormatContext<'ast>>) -> FormatResult<()> {
-        if let NodeLevel::Expression(Some(group_id)) = f.context().node_level() {
-            // If this content is enclosed by a group that adds the optional parentheses, then *disable*
-            // this group *except* if the optional parentheses are shown.
-            conditional_group(
-                &Arguments::from(&self.content),
-                Condition::if_group_breaks(group_id),
-            )
-            .fmt(f)
-        } else {
-            // Unconditionally group the content if it is not enclosed by an optional parentheses group.
-            group(&Arguments::from(&self.content)).fmt(f)
+        match f.context().node_level() {
+            NodeLevel::Expression(Some(parentheses_id)) => {
+                // If this content is enclosed by a group that adds the optional parentheses, then *disable*
+                // this group *except* if the optional parentheses are shown.
+                conditional_group(
+                    &Arguments::from(&self.content),
+                    Condition::if_group_breaks(parentheses_id),
+                )
+                .fmt(f)
+            }
+            NodeLevel::ParenthesizedExpression => {
+                // Unconditionally group the content if it is not enclosed by an optional parentheses group.
+                group(&Arguments::from(&self.content)).fmt(f)
+            }
+            NodeLevel::Expression(None) | NodeLevel::TopLevel | NodeLevel::CompoundStatement => {
+                Arguments::from(&self.content).fmt(f)
+            }
         }
     }
 }
diff --git a/crates/ruff_python_formatter/src/expression/string.rs b/crates/ruff_python_formatter/src/expression/string.rs
@@ -10,7 +10,9 @@ use ruff_formatter::{format_args, write, FormatError};
 use ruff_python_ast::str::is_implicit_concatenation;
 
 use crate::comments::{leading_comments, trailing_comments};
-use crate::expression::parentheses::in_parentheses_only_group;
+use crate::expression::parentheses::{
+    in_parentheses_only_group, in_parentheses_only_soft_line_break_or_space,
+};
 use crate::prelude::*;
 use crate::QuoteStyle;
 
@@ -64,7 +66,7 @@ impl Format<PyFormatContext<'_>> for FormatStringContinuation<'_> {
         // because this is a black preview style.
         let lexer = lex_starts_at(string_content, Mode::Expression, string_range.start());
 
-        let mut joiner = f.join_with(soft_line_break_or_space());
+        let mut joiner = f.join_with(in_parentheses_only_soft_line_break_or_space());
 
         for token in lexer {
             let (token, token_range) = match token {

diff --git a/crates/ruff_python_formatter/src/statement/stmt_class_def.rs b/crates/ruff_python_formatter/src/statement/stmt_class_def.rs
@@ -1,9 +1,9 @@
 use crate::comments::trailing_comments;
 
-use crate::expression::parentheses::Parentheses;
+use crate::expression::parentheses::{parenthesized, Parentheses};
 use crate::prelude::*;
 use crate::trivia::{SimpleTokenizer, TokenKind};
-use ruff_formatter::{format_args, write};
+use ruff_formatter::write;
 use ruff_text_size::TextRange;
 use rustpython_parser::ast::{Ranged, StmtClassDef};
 
@@ -32,16 +32,14 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
         write!(f, [text("class"), space(), name.format()])?;
 
         if !(bases.is_empty() && keywords.is_empty()) {
-            write!(
-                f,
-                [group(&format_args![
-                    text("("),
-                    soft_block_indent(&FormatInheritanceClause {
-                        class_definition: item
-                    }),
-                    text(")")
-                ])]
-            )?;
+            parenthesized(
+                "(",
+                &FormatInheritanceClause {
+                    class_definition: item,
+                },
+                ")",
+            )
+            .fmt(f)?;
         }
 
         let comments = f.context().comments().clone();

diff --git a/crates/ruff_python_formatter/src/statement/stmt_expr.rs b/crates/ruff_python_formatter/src/statement/stmt_expr.rs
@@ -1,4 +1,5 @@
-use rustpython_parser::ast::StmtExpr;
+use rustpython_parser::ast;
+use rustpython_parser::ast::{Expr, Operator, StmtExpr};
 
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
@@ -12,6 +13,25 @@ impl FormatNodeRule<StmtExpr> for FormatStmtExpr {
     fn fmt_fields(&self, item: &StmtExpr, f: &mut PyFormatter) -> FormatResult<()> {
         let StmtExpr { value, .. } = item;
 
-        maybe_parenthesize_expression(value, item, Parenthesize::Optional).fmt(f)
+        if is_arithmetic_like(value) {
+            maybe_parenthesize_expression(value, item, Parenthesize::Optional).fmt(f)
+        } else {
+            value.format().fmt(f)
+        }
     }
 }
+
+const fn is_arithmetic_like(expression: &Expr) -> bool {
+    matches!(
+        expression,
+        Expr::BinOp(ast::ExprBinOp {
+            op: Operator::BitOr
+                | Operator::BitXor
+                | Operator::LShift
+                | Operator::RShift
+                | Operator::Add
+                | Operator::Sub,
+            ..
+        })
+    )
+}
diff --git a/...n_formatter/tests/snapshots/black_compatibility@simple_cases__remove_await_parens.py.snap b/...n_formatter/tests/snapshots/black_compatibility@simple_cases__remove_await_parens.py.snap
@@ -114,8 +114,8 @@ async def main():
  # Check comments
  async def main():
 -    await asyncio.sleep(1)  # Hello
-+    (
-+        await # Hello
++    await (
++        # Hello
 +        asyncio.sleep(1)
 +    )
  
@@ -191,8 +191,8 @@ async def main():
 
 # Check comments
 async def main():
-    (
-        await # Hello
+    await (
+        # Hello
         asyncio.sleep(1)
     )
 

diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__compare.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__compare.py.snap
@@ -93,22 +93,12 @@ a not in b
 
 a < b > c == d
 
-(
-    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
-    < bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
-    > ccccccccccccccccccccccccccccc
-    == ddddddddddddddddddddd
-)
+aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa < bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb > ccccccccccccccccccccccccccccc == ddddddddddddddddddddd
 
-(
-    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
-    < [
-        bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb,
-        ff,
-    ]
-    < [ccccccccccccccccccccccccccccc, dddd]
-    < ddddddddddddddddddddddddddddddddddddddddddd
-)
+aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa < [
+    bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb,
+    ff,
+] < [ccccccccccccccccccccccccccccc, dddd] < ddddddddddddddddddddddddddddddddddddddddddd
 
 return 1 == 2 and (
     name,