Format target: annotation = value? expressions (#5661)

crates/ruff_python_formatter/src/builders.rs

@@ -6,20 +6,20 @@ use ruff_text_size::TextSize;
 use rustpython_parser::ast::Ranged;
 
 /// Adds parentheses and indents `content` if it doesn't fit on a line.
-pub(crate) fn optional_parentheses<'ast, T>(content: &T) -> OptionalParentheses<'_, 'ast>
+pub(crate) fn parenthesize_if_expands<'ast, T>(content: &T) -> ParenthesizeIfExpands<'_, 'ast>
 where
     T: Format<PyFormatContext<'ast>>,
 {
-    OptionalParentheses {
+    ParenthesizeIfExpands {
         inner: Argument::new(content),
     }
 }
 
-pub(crate) struct OptionalParentheses<'a, 'ast> {
+pub(crate) struct ParenthesizeIfExpands<'a, 'ast> {
     inner: Argument<'a, PyFormatContext<'ast>>,
 }
 
-impl<'ast> Format<PyFormatContext<'ast>> for OptionalParentheses<'_, 'ast> {
+impl<'ast> Format<PyFormatContext<'ast>> for ParenthesizeIfExpands<'_, 'ast> {
     fn fmt(&self, f: &mut Formatter<PyFormatContext<'ast>>) -> FormatResult<()> {
         let saved_level = f.context().node_level();
 

crates/ruff_python_formatter/src/expression/expr_subscript.rs

@@ -1,4 +1,4 @@
-use rustpython_parser::ast::ExprSubscript;
+use rustpython_parser::ast::{Expr, ExprSubscript};
 
 use ruff_formatter::{format_args, write};
 use ruff_python_ast::node::AstNode;
@@ -6,8 +6,10 @@ use ruff_python_ast::node::AstNode;
 use crate::comments::trailing_comments;
 use crate::context::NodeLevel;
 use crate::context::PyFormatContext;
+use crate::expression::expr_tuple::TupleParentheses;
 use crate::expression::parentheses::{
-    default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
+    default_expression_needs_parentheses, in_parentheses_only_group, NeedsParentheses, Parentheses,
+    Parenthesize,
 };
 use crate::prelude::*;
 use crate::FormatNodeRule;
@@ -42,12 +44,31 @@ impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
             value.format().fmt(f)?;
         }
 
+        let format_slice = format_with(|f: &mut PyFormatter| {
+            let saved_level = f.context().node_level();
+            f.context_mut()
+                .set_node_level(NodeLevel::ParenthesizedExpression);
+
+            let result = if let Expr::Tuple(tuple) = slice.as_ref() {
+                tuple
+                    .format()
+                    .with_options(TupleParentheses::Subscript)
+                    .fmt(f)
+            } else {
+                slice.format().fmt(f)
+            };
+
+            f.context_mut().set_node_level(saved_level);
+
+            result
+        });
+
         write!(
             f,
-            [group(&format_args![
+            [in_parentheses_only_group(&format_args![
                 text("["),
                 trailing_comments(dangling_comments),
-                soft_block_indent(&slice.format()),
+                soft_block_indent(&format_slice),
                 text("]")
             ])]
         )

crates/ruff_python_formatter/src/expression/expr_tuple.rs

@@ -1,4 +1,4 @@
-use crate::builders::optional_parentheses;
+use crate::builders::parenthesize_if_expands;
 use crate::comments::{dangling_comments, CommentLinePosition};
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, parenthesized, NeedsParentheses, Parentheses,
@@ -17,6 +17,11 @@ pub enum TupleParentheses {
     Default,
     /// Effectively `Some(Parentheses)` in `Option<Parentheses>`
     Expr(Parentheses),
+
+    /// Black omits parentheses for tuples inside of subscripts except if the tuple is parenthesized
+    /// in the source code.
+    Subscript,
+
     /// Handle the special case where we remove parentheses even if they were initially present
     ///
     /// Normally, black keeps parentheses, but in the case of loops it formats
@@ -86,21 +91,32 @@ impl FormatNodeRule<ExprTuple> for FormatExprTuple {
                     ])]
                 )
             }
-            [single] => {
-                // A single element tuple always needs parentheses and a trailing comma
-                parenthesized("(", &format_args![single.format(), &text(",")], ")").fmt(f)
-            }
+            [single] => match self.parentheses {
+                TupleParentheses::Subscript
+                    if !is_parenthesized(*range, elts, f.context().source()) =>
+                {
+                    write!(f, [single.format(), text(",")])
+                }
+                _ =>
+                // A single element tuple always needs parentheses and a trailing comma, except when inside of a subscript
+                {
+                    parenthesized("(", &format_args![single.format(), text(",")], ")").fmt(f)
+                }
+            },
             // If the tuple has parentheses, we generally want to keep them. The exception are for
             // loops, see `TupleParentheses::StripInsideForLoop` doc comment.
             //
             // Unlike other expression parentheses, tuple parentheses are part of the range of the
             // tuple itself.
-            elts if is_parenthesized(*range, elts, f)
+            elts if is_parenthesized(*range, elts, f.context().source())
                 && self.parentheses != TupleParentheses::StripInsideForLoop =>
             {
                 parenthesized("(", &ExprSequence::new(elts), ")").fmt(f)
             }
-            elts => optional_parentheses(&ExprSequence::new(elts)).fmt(f),
+            elts => match self.parentheses {
+                TupleParentheses::Subscript => group(&ExprSequence::new(elts)).fmt(f),
+                _ => parenthesize_if_expands(&ExprSequence::new(elts)).fmt(f),
+            },
         }
     }
 
@@ -141,15 +157,9 @@ impl NeedsParentheses for ExprTuple {
 }
 
 /// Check if a tuple has already had parentheses in the input
-fn is_parenthesized(
-    tuple_range: TextRange,
-    elts: &[Expr],
-    f: &mut Formatter<PyFormatContext<'_>>,
-) -> bool {
+fn is_parenthesized(tuple_range: TextRange, elts: &[Expr], source: &str) -> bool {
     let parentheses = '(';
-    let first_char = &f.context().source()[usize::from(tuple_range.start())..]
-        .chars()
-        .next();
+    let first_char = &source[usize::from(tuple_range.start())..].chars().next();
     let Some(first_char) = first_char else {
         return false;
     };

crates/ruff_python_formatter/src/expression/mod.rs

@@ -2,17 +2,16 @@ use rustpython_parser::ast;
 use rustpython_parser::ast::{Expr, Operator};
 use std::cmp::Ordering;
 
-use crate::builders::optional_parentheses;
-use ruff_formatter::{
-    format_args, FormatOwnedWithRule, FormatRefWithRule, FormatRule, FormatRuleWithOptions,
-};
+use crate::builders::parenthesize_if_expands;
+use ruff_formatter::{FormatOwnedWithRule, FormatRefWithRule, FormatRule, FormatRuleWithOptions};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::visitor::preorder::{walk_expr, PreorderVisitor};
 
 use crate::context::NodeLevel;
 use crate::expression::expr_tuple::TupleParentheses;
 use crate::expression::parentheses::{
-    is_expression_parenthesized, parenthesized, NeedsParentheses, Parentheses, Parenthesize,
+    is_expression_parenthesized, optional_parentheses, parenthesized, NeedsParentheses,
+    Parentheses, Parenthesize,
 };
 use crate::expression::string::StringLayout;
 use crate::prelude::*;
@@ -106,37 +105,9 @@ impl FormatRule<Expr, PyFormatContext<'_>> for FormatExpr {
             // Add optional parentheses. Ignore if the item renders parentheses itself.
             Parentheses::Optional => {
                 if can_omit_optional_parentheses(item, f.context()) {
-                    let saved_level = f.context().node_level();
-
-                    // The group id is used as a condition in [`in_parentheses_only`] to create a conditional group
-                    // that is only active if the optional parentheses group expands.
-                    let parens_id = f.group_id("optional_parentheses");
-
-                    f.context_mut()
-                        .set_node_level(NodeLevel::Expression(Some(parens_id)));
-
-                    // We can't use `soft_block_indent` here because that would always increment the indent,
-                    // even if the group does not break (the indent is not soft). This would result in
-                    // too deep indentations if a `parenthesized` group expands. Using `indent_if_group_breaks`
-                    // gives us the desired *soft* indentation that is only present if the optional parentheses
-                    // are shown.
-                    let result = group(&format_args![
-                        if_group_breaks(&text("(")),
-                        indent_if_group_breaks(
-                            &format_args![soft_line_break(), format_expr],
-                            parens_id
-                        ),
-                        soft_line_break(),
-                        if_group_breaks(&text(")"))
-                    ])
-                    .with_group_id(Some(parens_id))
-                    .fmt(f);
-
-                    f.context_mut().set_node_level(saved_level);
-
-                    result
-                } else {
                     optional_parentheses(&format_expr).fmt(f)
+                } else {
+                    parenthesize_if_expands(&format_expr).fmt(f)
                 }
             }
             Parentheses::Custom | Parentheses::Never => {

crates/ruff_python_formatter/src/expression/parentheses.rs

@@ -178,6 +178,57 @@ impl<'ast> Format<PyFormatContext<'ast>> for FormatParenthesized<'_, 'ast> {
     }
 }
 
+/// Wraps an expression in parentheses only if it still does not fit after expanding all expressions that start or end with
+/// a parentheses (`()`, `[]`, `{}`).
+pub(crate) fn optional_parentheses<'content, 'ast, Content>(
+    content: &'content Content,
+) -> OptionalParentheses<'content, 'ast>
+where
+    Content: Format<PyFormatContext<'ast>>,
+{
+    OptionalParentheses {
+        content: Argument::new(content),
+    }
+}
+
+pub(crate) struct OptionalParentheses<'content, 'ast> {
+    content: Argument<'content, PyFormatContext<'ast>>,
+}
+
+impl<'ast> Format<PyFormatContext<'ast>> for OptionalParentheses<'_, 'ast> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'ast>>) -> FormatResult<()> {
+        let saved_level = f.context().node_level();
+
+        // The group id is used as a condition in [`in_parentheses_only`] to create a conditional group
+        // that is only active if the optional parentheses group expands.
+        let parens_id = f.group_id("optional_parentheses");
+
+        f.context_mut()
+            .set_node_level(NodeLevel::Expression(Some(parens_id)));
+
+        // We can't use `soft_block_indent` here because that would always increment the indent,
+        // even if the group does not break (the indent is not soft). This would result in
+        // too deep indentations if a `parenthesized` group expands. Using `indent_if_group_breaks`
+        // gives us the desired *soft* indentation that is only present if the optional parentheses
+        // are shown.
+        let result = group(&format_args![
+            if_group_breaks(&text("(")),
+            indent_if_group_breaks(
+                &format_args![soft_line_break(), Arguments::from(&self.content)],
+                parens_id
+            ),
+            soft_line_break(),
+            if_group_breaks(&text(")"))
+        ])
+        .with_group_id(Some(parens_id))
+        .fmt(f);
+
+        f.context_mut().set_node_level(saved_level);
+
+        result
+    }
+}
+
 /// 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>(

crates/ruff_python_formatter/src/expression/string.rs

@@ -1,4 +1,4 @@
-use crate::builders::optional_parentheses;
+use crate::builders::parenthesize_if_expands;
 use crate::comments::{leading_comments, trailing_comments};
 use crate::expression::parentheses::Parentheses;
 use crate::prelude::*;
@@ -48,7 +48,7 @@ impl<'a> Format<PyFormatContext<'_>> for FormatString<'a> {
             let format_continuation = FormatStringContinuation::new(self.constant, self.layout);
 
             if let StringLayout::Default(Some(Parentheses::Custom)) = self.layout {
-                optional_parentheses(&format_continuation).fmt(f)
+                parenthesize_if_expands(&format_continuation).fmt(f)
             } else {
                 format_continuation.fmt(f)
             }

crates/ruff_python_formatter/src/lib.rs

@@ -280,11 +280,9 @@ if True:
     #[test]
     fn quick_test() {
         let src = r#"
-if [
-    aaaaaa,
-    BBBB,ccccccccc,ddddddd,eeeeeeeeee,ffffff
-] & bbbbbbbbbbbbbbbbbbddddddddddddddddddddddddddddbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb:
-    ...
+def foo() -> tuple[int, int, int,]:
+    return 2
+
 "#;
         // Tokenize once
         let mut tokens = Vec::new();

crates/ruff_python_formatter/src/statement/stmt_ann_assign.rs

@@ -1,5 +1,6 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
+use crate::prelude::*;
+use crate::FormatNodeRule;
+use ruff_formatter::write;
 use rustpython_parser::ast::StmtAnnAssign;
 
 #[derive(Default)]
@@ -7,6 +8,23 @@ pub struct FormatStmtAnnAssign;
 
 impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
     fn fmt_fields(&self, item: &StmtAnnAssign, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented(item)])
+        let StmtAnnAssign {
+            range: _,
+            target,
+            annotation,
+            value,
+            simple: _,
+        } = item;
+
+        write!(
+            f,
+            [target.format(), text(":"), space(), annotation.format()]
+        )?;
+
+        if let Some(value) = value {
+            write!(f, [space(), text("="), space(), value.format()])?;
+        }
+
+        Ok(())
     }
 }

crates/ruff_python_formatter/src/statement/stmt_assign.rs

@@ -1,12 +1,11 @@
-use crate::context::PyFormatContext;
-use crate::expression::parentheses::Parenthesize;
-use crate::{AsFormat, FormatNodeRule, PyFormatter};
-use ruff_formatter::formatter::Formatter;
-use ruff_formatter::prelude::{space, text};
-use ruff_formatter::{write, Buffer, Format, FormatResult};
-use rustpython_parser::ast::Expr;
 use rustpython_parser::ast::StmtAssign;
 
+use ruff_formatter::write;
+
+use crate::expression::parentheses::Parenthesize;
+use crate::prelude::*;
+use crate::FormatNodeRule;
+
 // Note: This currently does wrap but not the black way so the types below likely need to be
 // replaced entirely
 //
@@ -22,32 +21,11 @@ impl FormatNodeRule<StmtAssign> for FormatStmtAssign {
             value,
             type_comment: _,
         } = item;
-        write!(
-            f,
-            [
-                LhsAssignList::new(targets),
-                value.format().with_options(Parenthesize::IfBreaks)
-            ]
-        )
-    }
-}
 
-#[derive(Debug)]
-struct LhsAssignList<'a> {
-    lhs_assign_list: &'a [Expr],
-}
-
-impl<'a> LhsAssignList<'a> {
-    const fn new(lhs_assign_list: &'a [Expr]) -> Self {
-        Self { lhs_assign_list }
-    }
-}
-
-impl Format<PyFormatContext<'_>> for LhsAssignList<'_> {
-    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
-        for element in self.lhs_assign_list {
-            write!(f, [&element.format(), space(), text("="), space(),])?;
+        for target in targets {
+            write!(f, [target.format(), space(), text("="), space()])?;
         }
-        Ok(())
+
+        write!(f, [value.format().with_options(Parenthesize::IfBreaks)])
     }
 }

crates/ruff_python_formatter/src/statement/stmt_delete.rs

@@ -1,4 +1,4 @@
-use crate::builders::{optional_parentheses, PyFormatterExtensions};
+use crate::builders::{parenthesize_if_expands, PyFormatterExtensions};
 use crate::comments::dangling_node_comments;
 use crate::expression::parentheses::Parenthesize;
 use crate::{AsFormat, FormatNodeRule, PyFormatter};
@@ -36,7 +36,7 @@ impl FormatNodeRule<StmtDelete> for FormatStmtDelete {
             }
             targets => {
                 let item = format_with(|f| f.join_comma_separated().nodes(targets.iter()).finish());
-                optional_parentheses(&item).fmt(f)
+                parenthesize_if_expands(&item).fmt(f)
             }
         }
     }

crates/ruff_python_formatter/src/statement/stmt_function_def.rs

@@ -1,6 +1,6 @@
 use crate::comments::{leading_comments, trailing_comments};
 use crate::context::NodeLevel;
-use crate::expression::parentheses::Parenthesize;
+use crate::expression::parentheses::{optional_parentheses, Parenthesize};
 use crate::prelude::*;
 use crate::trivia::{lines_after, skip_trailing_trivia};
 use crate::FormatNodeRule;
@@ -97,9 +97,9 @@ impl FormatRule<AnyFunctionDefinition<'_>, PyFormatContext<'_>> for FormatAnyFun
                     space(),
                     text("->"),
                     space(),
-                    return_annotation
-                        .format()
-                        .with_options(Parenthesize::IfBreaks)
+                    optional_parentheses(
+                        &return_annotation.format().with_options(Parenthesize::Never)
+                    )
                 ]
             )?;
         }

crates/ruff_python_formatter/src/statement/stmt_import_from.rs

@@ -1,4 +1,4 @@
-use crate::builders::{optional_parentheses, PyFormatterExtensions};
+use crate::builders::{parenthesize_if_expands, PyFormatterExtensions};
 use crate::{AsFormat, FormatNodeRule, PyFormatter};
 use ruff_formatter::prelude::{dynamic_text, format_with, space, text};
 use ruff_formatter::{write, Buffer, Format, FormatResult};
@@ -43,6 +43,6 @@ impl FormatNodeRule<StmtImportFrom> for FormatStmtImportFrom {
                 .entries(names.iter().map(|name| (name, name.format())))
                 .finish()
         });
-        optional_parentheses(&names).fmt(f)
+        parenthesize_if_expands(&names).fmt(f)
     }
 }

crates/ruff_python_formatter/src/statement/stmt_with.rs

@@ -3,7 +3,7 @@ use ruff_python_ast::node::AnyNodeRef;
 use ruff_text_size::TextRange;
 use rustpython_parser::ast::{Ranged, StmtAsyncWith, StmtWith, Suite, WithItem};
 
-use crate::builders::optional_parentheses;
+use crate::builders::parenthesize_if_expands;
 use crate::comments::trailing_comments;
 use crate::prelude::*;
 use crate::FormatNodeRule;
@@ -80,7 +80,7 @@ impl Format<PyFormatContext<'_>> for AnyStatementWith<'_> {
             [
                 text("with"),
                 space(),
-                group(&optional_parentheses(&joined_items)),
+                group(&parenthesize_if_expands(&joined_items)),
                 text(":"),
                 trailing_comments(dangling_comments),
                 block_indent(&self.body().format())