Fix unnecessary parentheses in indexed assignment RHS

CHANGES.md

@@ -21,6 +21,8 @@
 
 <!-- Changes that affect Black's preview style -->
 
+- Fix unnecessary parentheses around short RHS expressions in indexed assignments like
+  `x[key] = expr` (#5095)
 - Prevent string merger from creating unsplittable long lines when a pragma comment
   (e.g. `# type: ignore`) follows the closing bracket (#5096)
 - Improve heuristics around whether blank lines should appear before, within and after

docs/the_black_code_style/future_style.md

@@ -31,6 +31,9 @@ Currently, the following features are included in the preview style:
 - `pyi_blank_line_before_decorated_class`: In `.pyi` stub files, enforce a blank line
   before a decorated class definition when it follows a function definition.
   ([see below](labels/pyi-blank-line-before-decorated-class))
+- `fix_unnecessary_parens_in_indexed_assignment`: Remove unnecessary parentheses around
+  the right-hand side of indexed assignments (e.g. `x[key] = expr`) when the expression
+  is short enough. ([see below](labels/fix-unnecessary-parens-indexed-assignment))
 
 (labels/wrap-comprehension-in)=
 
@@ -230,6 +233,33 @@ classes are handled:
   class Spam: ...
 ```
 
+(labels/fix-unnecessary-parens-indexed-assignment)=
+
+### Unnecessary parentheses in indexed assignments
+
+When an assignment target contains brackets (e.g. indexed access like `x[key] = expr`),
+previously, Black would incorrectly wrap the right-hand side expression in unnecessary
+parentheses when the line was too long. With this feature enabled, Black removes the
+unnecessary parentheses when the RHS expression fits on the tail line.
+
+For example:
+
+```python
+# Before
+dictionary_of_arrays["long_key_name_for_the_example"][
+    very_long_index_name, index_zero
+] = (10 - 5)
+```
+
+will be formatted to:
+
+```python
+# After (with --preview)
+dictionary_of_arrays["long_key_name_for_the_example"][
+    very_long_index_name, index_zero
+] = 10 - 5
+```
+
 ## Unstable style
 
 (labels/unstable-style)=

src/black/linegen.py

@@ -1077,7 +1077,7 @@ def _maybe_split_omitting_optional_parens(
         # in this case; attempting a split without them is a waste of time)
         and not line.is_import
         # and we can actually remove the parens
-        and can_omit_invisible_parens(rhs, mode.line_length)
+        and can_omit_invisible_parens(rhs, mode.line_length, mode)
     ):
         omit = {id(rhs.closing_bracket), *omit}
         try:
@@ -1180,6 +1180,24 @@ def _prefer_split_rhs_oop_over_rhs(
     if rhs_head_equal_count > 1 and rhs_head_equal_count > rhs_oop_head_equal_count:
         return False
 
+    # Don't prefer OOP if the split breaks inside a subscript access chain on the
+    # RHS. When rhs_oop.tail starts with `]` followed by more tokens (like `.attr`
+    # or `.method()`), and `=` is in rhs_oop.head (split is in the RHS, not the
+    # LHS), this means we split mid-chain — creating ugly output like
+    # `expr + obj[\n    idx\n].attr`. Prefer paren-wrapping instead.
+    # Gated behind the same preview flag as the indexed assignment fix in
+    # can_omit_invisible_parens, since this guard only matters when that path
+    # is active (stable mode is unaffected).
+    if (
+        Preview.fix_unnecessary_parens_in_indexed_assignment in mode
+        and rhs_oop.tail.leaves
+        and rhs_oop.tail.leaves[0].type == token.RSQB
+        and len(rhs_oop.tail.leaves) > 1
+        and rhs_oop.tail.leaves[1].type == token.DOT
+        and any(leaf.type == token.EQUAL for leaf in rhs_oop.head.leaves)
+    ):
+        return False
+
     has_closing_bracket_after_assign = False
     for leaf in reversed(rhs_oop.head.leaves):
         if leaf.type == token.EQUAL:

src/black/lines.py

@@ -1366,6 +1366,7 @@ def can_be_split(line: Line) -> bool:
 def can_omit_invisible_parens(
     rhs: RHSResult,
     line_length: int,
+    mode: Mode,
 ) -> bool:
     """Does `rhs.body` have a shape safe to reformat without optional parens around it?
 
@@ -1429,6 +1430,28 @@ def can_omit_invisible_parens(
         ):
             closing_bracket = leaf
 
+    # For assignment RHS where the LHS contains brackets (e.g. indexed
+    # assignments like `x[key] = expr`), allow omitting optional parens
+    # if the body fits on a single line or has brackets to further split on.
+    # This check must come before the delimiter_count > 1 early return below,
+    # which would otherwise reject expressions like `1 + 1 + 1 + ...`.
+    # The downstream _prefer_split_rhs_oop_over_rhs will make the final
+    # decision on whether the result is actually better.
+    if (
+        Preview.fix_unnecessary_parens_in_indexed_assignment in mode
+        and len(rhs.head.leaves) >= 2
+        and rhs.head.leaves[-2].type == token.EQUAL
+        and any(leaf.type in BRACKETS and leaf.value for leaf in rhs.head.leaves[:-2])
+    ):
+        body_length = 4 * line.depth
+        has_brackets = False
+        for _index, _leaf, leaf_length in line.enumerate_with_length():
+            body_length += leaf_length
+            if _leaf.type in OPENING_BRACKETS:
+                has_brackets = True
+        if has_brackets or body_length <= line_length:
+            return True
+
     bt = line.bracket_tracker
     if not bt.delimiters:
         # Without delimiters the optional parentheses are useless.

src/black/mode.py

@@ -258,6 +258,7 @@ class Preview(Enum):
     wrap_long_dict_values_in_parens = auto()
     fix_if_guard_explosion_in_case_statement = auto()
     pyi_overload_group_blank_lines = auto()
+    fix_unnecessary_parens_in_indexed_assignment = auto()
     pyi_blank_line_before_decorated_class = auto()
 
 

src/black/resources/black.schema.json

@@ -88,6 +88,7 @@
           "wrap_long_dict_values_in_parens",
           "fix_if_guard_explosion_in_case_statement",
           "pyi_overload_group_blank_lines",
+          "fix_unnecessary_parens_in_indexed_assignment",
           "pyi_blank_line_before_decorated_class"
         ]
       },

tests/data/cases/prefer_rhs_split.py

@@ -104,3 +104,15 @@
 ) = (
     cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
 ) = ddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd
+
+
+# Make sure that when the RHS contains a subscript access chain (attribute
+# access with subscript like obj.children[idx].attr), Black prefers wrapping
+# the RHS in parentheses rather than splitting at the subscript.
+some_node.children[1].prefix = (
+    some_node.children[0].prefix + some_node.children[1].prefix
+)
+
+another_node.children[idx].value = (
+    another_node.children[idx - 1].value + another_node.children[idx + 1].value
+)

tests/data/cases/prefer_rhs_split_reformatted.py

@@ -20,6 +20,11 @@
     )
 )
 
+# Regression: subscript access chain on RHS should not be split mid-chain
+some_node.children[1].prefix = some_node.children[0].prefix + some_node.children[1].prefix
+
+another_node.children[idx].value = another_node.children[idx - 1].value + another_node.children[idx + 1].value
+
 # output
 
 
@@ -55,3 +60,12 @@
         c=3,
     )
 )
+
+# Regression: subscript access chain on RHS should not be split mid-chain
+some_node.children[1].prefix = (
+    some_node.children[0].prefix + some_node.children[1].prefix
+)
+
+another_node.children[idx].value = (
+    another_node.children[idx - 1].value + another_node.children[idx + 1].value
+)

tests/data/cases/preview_prefer_rhs_split_indexed_assignment.py

@@ -0,0 +1,37 @@
+# flags: --preview
+
+# Indexed assignment with a short RHS expression should not get unnecessary parens.
+dictionary_of_arrays["long_key_name_for_the_example"][
+    very_long_index_name, index_zero
+] = 10 - 5
+
+# Unformatted input: the unnecessary parens should be removed.
+dictionary_of_arrays["long_key_name_for_the_example"][very_long_index_name, index_zero] = (10 - 5)
+
+# Longer RHS expressions that fit on the tail line should also lose parens.
+dictionary_of_arrays["long_key_name_for_the_example"][very_long_index_name, index_zero] = (1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1)
+
+# RHS with brackets (function call) should also lose unnecessary parens.
+dictionary_of_arrays["long_key_name_for_the_example"][very_long_index_name, index_zero] = (some_function(arg1, arg2))
+
+# output
+
+# Indexed assignment with a short RHS expression should not get unnecessary parens.
+dictionary_of_arrays["long_key_name_for_the_example"][
+    very_long_index_name, index_zero
+] = 10 - 5
+
+# Unformatted input: the unnecessary parens should be removed.
+dictionary_of_arrays["long_key_name_for_the_example"][
+    very_long_index_name, index_zero
+] = 10 - 5
+
+# Longer RHS expressions that fit on the tail line should also lose parens.
+dictionary_of_arrays["long_key_name_for_the_example"][
+    very_long_index_name, index_zero
+] = 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1
+
+# RHS with brackets (function call) should also lose unnecessary parens.
+dictionary_of_arrays["long_key_name_for_the_example"][
+    very_long_index_name, index_zero
+] = some_function(arg1, arg2)