Fix: Avoid parenthesizing subscript targets and values

Signed-off-by: Micha Reiser <[email protected]>

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/assignment_split_value_first.py

@@ -153,6 +153,23 @@
     function().b().c([1, 2, 3], arg1, [1, 2, 3], arg2, [1, 2, 3], arg3)
 )
 
+#######
+# Subscripts and non-fluent attribute chains
+a = aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa[
+    xxxxx
+].bbvbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb[
+    yyyyyyyyyy[aaaa]
+] = ccccccccccccccccccccccccccccccccccc["aaaaaaa"]
+
+a = aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa[
+    xxxxx
+].bbvbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb = ccccccccccccccccccccccccccccccccccc[
+    "aaaaaaa"
+]
+
+label_thresholds[label_id] = label_quantiles[label_id][
+    min(int(tolerance * num_thresholds), num_thresholds - 1)
+]
 
 #######
 # Test comment inlining

crates/ruff_python_formatter/src/statement/stmt_assign.rs

@@ -1,5 +1,7 @@
 use ruff_formatter::{format_args, write, FormatError};
-use ruff_python_ast::{AnyNodeRef, Expr, Operator, StmtAssign, TypeParams};
+use ruff_python_ast::{
+    AnyNodeRef, Expr, ExprAttribute, ExprCall, Operator, StmtAssign, TypeParams,
+};
 
 use crate::builders::parenthesize_if_expands;
 use crate::comments::{
@@ -9,7 +11,7 @@ use crate::context::{NodeLevel, WithNodeLevel};
 use crate::expression::parentheses::{
     is_expression_parenthesized, NeedsParentheses, OptionalParentheses, Parentheses, Parenthesize,
 };
-use crate::expression::{has_own_parentheses, maybe_parenthesize_expression};
+use crate::expression::{has_own_parentheses, has_parentheses, maybe_parenthesize_expression};
 use crate::prelude::*;
 use crate::preview::is_prefer_splitting_right_hand_side_of_assignments_enabled;
 use crate::statement::trailing_semicolon;
@@ -56,7 +58,7 @@ impl FormatNodeRule<StmtAssign> for FormatStmtAssign {
                 }
                 .fmt(f)?;
             }
-            // Avoid parenthesizing the value for single-target assignments that where the
+            // Avoid parenthesizing the value for single-target assignments where the
             // target has its own parentheses (list, dict, tuple, ...) and the target expands.
             else if has_target_own_parentheses(first, f.context())
                 && !is_expression_parenthesized(
@@ -193,14 +195,14 @@ impl Format<PyFormatContext<'_>> for FormatTargetWithEqualOperator<'_> {
             || f.context().comments().has_trailing(self.target)
         {
             self.target.format().fmt(f)?;
-        } else if has_target_own_parentheses(self.target, f.context()) {
+        } else if should_parenthesize_target(self.target, f.context()) {
+            parenthesize_if_expands(&self.target.format().with_options(Parentheses::Never))
+                .fmt(f)?;
+        } else {
             self.target
                 .format()
                 .with_options(Parentheses::Never)
                 .fmt(f)?;
-        } else {
-            parenthesize_if_expands(&self.target.format().with_options(Parentheses::Never))
-                .fmt(f)?;
         }
 
         write!(f, [space(), token("="), space()])
@@ -554,7 +556,9 @@ impl Format<PyFormatContext<'_>> for FormatStatementsLastExpression<'_> {
 
                 // For call expressions, prefer breaking after the call expression's opening parentheses
                 // over parenthesizing the entire call expression.
-                if value.is_call_expr() {
+                // For subscripts, try breaking the subscript first
+                // For attribute chains that contain any parenthesized value: Try expanding the parenthesized value first.
+                if value.is_call_expr() || value.is_subscript_expr() || value.is_attribute_expr() {
                     best_fitting![
                         format_flat,
                         // Avoid parenthesizing the call expression if the `(` fit on the line
@@ -681,11 +685,11 @@ impl Format<PyFormatContext<'_>> for AnyBeforeOperator<'_> {
                         .fmt(f)
                 }
                 // Never parenthesize targets that come with their own parentheses, e.g. don't parenthesize lists or dictionary literals.
-                else if has_target_own_parentheses(expression, f.context()) {
-                    expression.format().with_options(Parentheses::Never).fmt(f)
-                } else {
+                else if should_parenthesize_target(expression, f.context()) {
                     parenthesize_if_expands(&expression.format().with_options(Parentheses::Never))
                         .fmt(f)
+                } else {
+                    expression.format().with_options(Parentheses::Never).fmt(f)
                 }
             }
             // Never parenthesize type params
@@ -717,7 +721,7 @@ fn should_inline_comments(
     }
 }
 
-/// Tests whether an expression that for which comments shouldn't be inlined should use the best fit layout
+/// Tests whether an expression for which comments shouldn't be inlined should use the best fit layout
 fn should_non_inlineable_use_best_fit(
     expr: &Expr,
     parent: AnyNodeRef,
@@ -728,12 +732,32 @@ fn should_non_inlineable_use_best_fit(
             attribute.needs_parentheses(parent, context) == OptionalParentheses::BestFit
         }
         Expr::Call(call) => call.needs_parentheses(parent, context) == OptionalParentheses::BestFit,
+        Expr::Subscript(subscript) => {
+            subscript.needs_parentheses(parent, context) == OptionalParentheses::BestFit
+        }
         _ => false,
     }
 }
 
-/// Returns `true` for targets that should not be parenthesized if they split because their expanded
-/// layout comes with their own set of parentheses.
+/// Returns `true` for targets that have their own set of parentheses when they split,
+/// in which case we want to avoid parenthesizing the assigned value.
 pub(super) fn has_target_own_parentheses(target: &Expr, context: &PyFormatContext) -> bool {
     matches!(target, Expr::Tuple(_)) || has_own_parentheses(target, context).is_some()
 }
+
+pub(super) fn should_parenthesize_target(target: &Expr, context: &PyFormatContext) -> bool {
+    !(has_target_own_parentheses(target, context)
+        || is_attribute_with_parenthesized_value(target, context))
+}
+
+fn is_attribute_with_parenthesized_value(target: &Expr, context: &PyFormatContext) -> bool {
+    match target {
+        Expr::Attribute(ExprAttribute { value, .. }) => {
+            has_parentheses(value.as_ref(), context).is_some()
+                || is_attribute_with_parenthesized_value(value, context)
+        }
+        Expr::Subscript(_) => true,
+        Expr::Call(ExprCall { arguments, .. }) => !arguments.is_empty(),
+        _ => false,
+    }
+}

crates/ruff_python_formatter/tests/snapshots/format@statement__assignment_split_value_first.py.snap

@@ -159,6 +159,23 @@ this_is_a_ridiculously_long_name_and_nobody_in_their_right_mind_would_use = (
     function().b().c([1, 2, 3], arg1, [1, 2, 3], arg2, [1, 2, 3], arg3)
 )
 
+#######
+# Subscripts and non-fluent attribute chains
+a = aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa[
+    xxxxx
+].bbvbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb[
+    yyyyyyyyyy[aaaa]
+] = ccccccccccccccccccccccccccccccccccc["aaaaaaa"]
+
+a = aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa[
+    xxxxx
+].bbvbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb = ccccccccccccccccccccccccccccccccccc[
+    "aaaaaaa"
+]
+
+label_thresholds[label_id] = label_quantiles[label_id][
+    min(int(tolerance * num_thresholds), num_thresholds - 1)
+]
 
 #######
 # Test comment inlining
@@ -394,6 +411,23 @@ this_is_a_ridiculously_long_name_and_nobody_in_their_right_mind_would_use = (
     function().b().c([1, 2, 3], arg1, [1, 2, 3], arg2, [1, 2, 3], arg3)
 )
 
+#######
+# Subscripts and non-fluent attribute chains
+a = aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa[
+    xxxxx
+].bbvbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb[
+    yyyyyyyyyy[aaaa]
+] = ccccccccccccccccccccccccccccccccccc["aaaaaaa"]
+
+a = aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa[
+    xxxxx
+].bbvbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb = ccccccccccccccccccccccccccccccccccc[
+    "aaaaaaa"
+]
+
+label_thresholds[label_id] = label_quantiles[label_id][
+    min(int(tolerance * num_thresholds), num_thresholds - 1)
+]
 
 #######
 # Test comment inlining