Format target: annotation = value? expressions

crates/ruff_python_formatter/resources/test/fixtures/black/simple_cases/skip_magic_trailing_comma.options.json

@@ -0,0 +1,3 @@
+{
+    "magic_trailing_comma": "ignore"
+}

crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/annotated_assign.py

@@ -0,0 +1,13 @@
+a: string
+
+b: string = "test"
+
+b: list[
+    string,
+    int
+] = [1, 2]
+
+b: list[
+    string,
+    int,
+] = [1, 2]

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,13 +1,15 @@
-use rustpython_parser::ast::ExprSubscript;
+use rustpython_parser::ast::{Expr, ExprSubscript};
 
 use ruff_formatter::{format_args, write};
 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("]")
             ])]
         )
@@ -69,6 +90,10 @@ impl NeedsParentheses for ExprSubscript {
         parenthesize: Parenthesize,
         context: &PyFormatContext,
     ) -> Parentheses {
+        // TODO: Issue, returns `Never` here but parentheses are necessary in return type positions.
+        // We don't have the position information available here which sux.
+        // Rename `optional_parentheses` to `parenthesize_if_breaks`
+        // Create a new optional parentheses function that does the whole magic stuff.
         match default_expression_needs_parentheses(self.into(), parenthesize, context) {
             Parentheses::Optional => Parentheses::Never,
             parentheses => parentheses,

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::expression::ExpressionRef;
 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,30 +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();
-
-                    let parens_id = f.group_id("optional_parentheses");
-
-                    f.context_mut()
-                        .set_node_level(NodeLevel::Expression(Some(parens_id)));
-
-                    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

@@ -170,6 +170,50 @@ 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();
+
+        let parens_id = f.group_id("optional_parentheses");
+
+        f.context_mut()
+            .set_node_level(NodeLevel::Expression(Some(parens_id)));
+
+        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,12 +1,30 @@
-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)]
 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(())
     }
 }