fix: Use BestFit layout for subscript (#7409)

crates/ruff_python_formatter/src/expression/expr_call.rs

@@ -4,7 +4,7 @@ use ruff_python_ast::{Expr, ExprCall};
 
 use crate::comments::{dangling_comments, SourceComment};
 use crate::expression::parentheses::{
-    is_expression_parenthesized, NeedsParentheses, OptionalParentheses,
+    is_expression_parenthesized, NeedsParentheses, OptionalParentheses, Parentheses,
 };
 use crate::expression::CallChainLayout;
 use crate::prelude::*;
@@ -36,13 +36,17 @@ impl FormatNodeRule<ExprCall> for FormatExprCall {
 
         let call_chain_layout = self.call_chain_layout.apply_in_node(item, f);
 
-        let fmt_func = format_with(|f| {
+        let fmt_func = format_with(|f: &mut PyFormatter| {
             // Format the function expression.
-            match func.as_ref() {
-                Expr::Attribute(expr) => expr.format().with_options(call_chain_layout).fmt(f),
-                Expr::Call(expr) => expr.format().with_options(call_chain_layout).fmt(f),
-                Expr::Subscript(expr) => expr.format().with_options(call_chain_layout).fmt(f),
-                _ => func.format().fmt(f),
+            if is_expression_parenthesized(func.into(), f.context().source()) {
+                func.format().with_options(Parentheses::Always).fmt(f)
+            } else {
+                match func.as_ref() {
+                    Expr::Attribute(expr) => expr.format().with_options(call_chain_layout).fmt(f),
+                    Expr::Call(expr) => expr.format().with_options(call_chain_layout).fmt(f),
+                    Expr::Subscript(expr) => expr.format().with_options(call_chain_layout).fmt(f),
+                    _ => func.format().with_options(Parentheses::Never).fmt(f),
+                }
             }?;
 
             // Format comments between the function and its arguments.

crates/ruff_python_formatter/src/expression/expr_subscript.rs

@@ -5,7 +5,7 @@ use ruff_python_ast::{Expr, ExprSubscript};
 use crate::comments::SourceComment;
 use crate::expression::expr_tuple::TupleParentheses;
 use crate::expression::parentheses::{
-    is_expression_parenthesized, parenthesized, NeedsParentheses, OptionalParentheses,
+    is_expression_parenthesized, parenthesized, NeedsParentheses, OptionalParentheses, Parentheses,
 };
 use crate::expression::CallChainLayout;
 use crate::prelude::*;
@@ -42,13 +42,17 @@ impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
             "A subscript expression can only have a single dangling comment, the one after the bracket"
         );
 
-        let format_inner = format_with(|f| {
-            match value.as_ref() {
-                Expr::Attribute(expr) => expr.format().with_options(call_chain_layout).fmt(f)?,
-                Expr::Call(expr) => expr.format().with_options(call_chain_layout).fmt(f)?,
-                Expr::Subscript(expr) => expr.format().with_options(call_chain_layout).fmt(f)?,
-                _ => value.format().fmt(f)?,
-            }
+        let format_inner = format_with(|f: &mut PyFormatter| {
+            if is_expression_parenthesized(value.into(), f.context().source()) {
+                value.format().with_options(Parentheses::Always).fmt(f)
+            } else {
+                match value.as_ref() {
+                    Expr::Attribute(expr) => expr.format().with_options(call_chain_layout).fmt(f),
+                    Expr::Call(expr) => expr.format().with_options(call_chain_layout).fmt(f),
+                    Expr::Subscript(expr) => expr.format().with_options(call_chain_layout).fmt(f),
+                    _ => value.format().with_options(Parentheses::Never).fmt(f),
+                }
+            }?;
 
             let format_slice = format_with(|f: &mut PyFormatter| {
                 if let Expr::Tuple(tuple) = slice.as_ref() {
@@ -85,7 +89,7 @@ impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
 impl NeedsParentheses for ExprSubscript {
     fn needs_parentheses(
         &self,
-        _parent: AnyNodeRef,
+        parent: AnyNodeRef,
         context: &PyFormatContext,
     ) -> OptionalParentheses {
         {
@@ -104,7 +108,21 @@ impl NeedsParentheses for ExprSubscript {
                 OptionalParentheses::Never
             } else {
                 match self.value.needs_parentheses(self.into(), context) {
-                    OptionalParentheses::BestFit => OptionalParentheses::Never,
+                    OptionalParentheses::BestFit => {
+                        if parent.as_stmt_function_def().is_some_and(|function_def| {
+                            function_def
+                                .returns
+                                .as_deref()
+                                .and_then(Expr::as_subscript_expr)
+                                == Some(self)
+                        }) {
+                            // Don't use the best fitting layout for return type annotation because it results in the
+                            // return type expanding before the parameters.
+                            OptionalParentheses::Never
+                        } else {
+                            OptionalParentheses::BestFit
+                        }
+                    }
                     parentheses => parentheses,
                 }
             }

crates/ruff_python_formatter/src/lib.rs

@@ -217,12 +217,14 @@ if True:
     #[test]
     fn quick_test() {
         let src = r#"
-(header.timecnt * 5  # Transition times and types
-  + header.typecnt * 6  # Local time type records
-  + header.charcnt  # Time zone designations
-  + header.leapcnt * 8  # Leap second records
-  + header.isstdcnt  # Standard/wall indicators
-  + header.isutcnt)  # UT/local indicators
+def main() -> None:
+    if True:
+        some_very_long_variable_name_abcdefghijk = Foo()
+        some_very_long_variable_name_abcdefghijk = some_very_long_variable_name_abcdefghijk[
+            some_very_long_variable_name_abcdefghijk.some_very_long_attribute_name
+            == "This is a very long string abcdefghijk"
+        ]
+
 "#;
         // Tokenize once
         let mut tokens = Vec::new();