Automatic management of parentheses in assignments

Fixes #140

Note: this is an evolution but the end result needs to be different.  See
cantfit.py for some good examples on bad formatting caused by this change.

README.md

@@ -327,6 +327,11 @@ interesting cases:
 - `for (...) in (...):`
 - `assert (...), (...)`
 - `from X import (...)`
+- assignments like:
+  - `target = (...)`
+  - `target: type = (...)`
+  - `some, *un, packing = (...)`
+  - `augmented += (...)`
 
 In those cases, parentheses are removed when the entire statement fits
 in one line, or if the inner expression doesn't have any delimiters to
@@ -540,6 +545,9 @@ More details can be found in [CONTRIBUTING](CONTRIBUTING.md).
 
 * slices are now formatted according to PEP 8 (#178)
 
+* parentheses are now also managed automatically on the right-hand side
+  of assignments and return statements (#140)
+
 * math operators now use their respective priorities for delimiting multiline
   expressions (#148)
 

black.py

@@ -599,6 +599,22 @@ def show(cls, code: str) -> None:
     syms.term,
     syms.power,
 }
+ASSIGNMENTS = {
+    "=",
+    "+=",
+    "-=",
+    "*=",
+    "@=",
+    "/=",
+    "%=",
+    "&=",
+    "|=",
+    "^=",
+    "<<=",
+    ">>=",
+    "**=",
+    "//=",
+}
 COMPREHENSION_PRIORITY = 20
 COMMA_PRIORITY = 18
 TERNARY_PRIORITY = 16
@@ -994,8 +1010,9 @@ def is_complex_subscript(self, leaf: Leaf) -> bool:
             and subscript_start.type == syms.subscriptlist
         ):
             subscript_start = child_towards(subscript_start, leaf)
-        return subscript_start is not None and any(
+        return (
-            n.type in TEST_DESCENDANTS for n in subscript_start.pre_order()
+            subscript_start is not None
+            and any(n.type in TEST_DESCENDANTS for n in subscript_start.pre_order())
         )
 
     def __str__(self) -> str:
@@ -1252,7 +1269,7 @@ def visit_stmt(
         """Visit a statement.
 
         This implementation is shared for `if`, `while`, `for`, `try`, `except`,
-        `def`, `with`, `class`, and `assert`.
+        `def`, `with`, `class`, `assert` and assignments.
 
         The relevant Python language `keywords` for a given statement will be
         NAME leaves within it. This methods puts those on a separate line.
@@ -1368,6 +1385,8 @@ def __attrs_post_init__(self) -> None:
         self.visit_with_stmt = partial(v, keywords={"with"}, parens=Ø)
         self.visit_funcdef = partial(v, keywords={"def"}, parens=Ø)
         self.visit_classdef = partial(v, keywords={"class"}, parens=Ø)
+        self.visit_expr_stmt = partial(v, keywords=Ø, parens=ASSIGNMENTS)
+        self.visit_return_stmt = partial(v, keywords={"return"}, parens={"return"})
         self.visit_async_funcdef = self.visit_async_stmt
         self.visit_decorated = self.visit_decorators
 

tests/cantfit.py

@@ -25,3 +25,43 @@
     "eggs with spam and eggs and spam with eggs with spam and eggs and spam with eggs with spam and eggs and spam with eggs",
     this_is_a_ridiculously_long_name_and_nobody_in_their_right_mind_would_use_one_like_it=0,
 )
+
+
+# output
+
+
+# long variable name
+this_is_a_ridiculously_long_name_and_nobody_in_their_right_mind_would_use_one_like_it = (
+    0
+)
+this_is_a_ridiculously_long_name_and_nobody_in_their_right_mind_would_use_one_like_it = (
+    1
+)  # with a comment
+this_is_a_ridiculously_long_name_and_nobody_in_their_right_mind_would_use_one_like_it = [
+    1, 2, 3
+]
+this_is_a_ridiculously_long_name_and_nobody_in_their_right_mind_would_use_one_like_it = (
+    function()
+)
+this_is_a_ridiculously_long_name_and_nobody_in_their_right_mind_would_use_one_like_it = function(
+    arg1, arg2, arg3
+)
+this_is_a_ridiculously_long_name_and_nobody_in_their_right_mind_would_use_one_like_it = function(
+    [1, 2, 3], arg1, [1, 2, 3], arg2, [1, 2, 3], arg3
+)
+# long function name
+normal_name = (
+    but_the_function_name_is_now_ridiculously_long_and_it_is_still_super_annoying()
+)
+normal_name = but_the_function_name_is_now_ridiculously_long_and_it_is_still_super_annoying(
+    arg1, arg2, arg3
+)
+normal_name = but_the_function_name_is_now_ridiculously_long_and_it_is_still_super_annoying(
+    [1, 2, 3], arg1, [1, 2, 3], arg2, [1, 2, 3], arg3
+)
+# long arguments
+normal_name = normal_function_name(
+    "but with super long string arguments that on their own exceed the line limit so there's no way it can ever fit",
+    "eggs with spam and eggs and spam with eggs with spam and eggs and spam with eggs with spam and eggs and spam with eggs",
+    this_is_a_ridiculously_long_name_and_nobody_in_their_right_mind_would_use_one_like_it=0,
+)

tests/expression.diff

@@ -27,11 +27,11 @@
 -foo = (lambda port_id, ignore_missing: {"port1": port1_resource, "port2": port2_resource}[port_id])
 +lambda a, b, c=True, *, d=(1 << v2), e="str": a
 +lambda a, b, c=True, *vararg, d=(v1 << 2), e="str", **kwargs: a + b
-+foo = (
++foo = lambda port_id, ignore_missing: {
-+    lambda port_id, ignore_missing: {"port1": port1_resource, "port2": port2_resource}[
++    "port1": port1_resource, "port2": port2_resource
-+        port_id
++}[
-+    ]
++    port_id
-+)
++]
  1 if True else 2
  str or None if True else str or bytes or None
  (str or None) if True else (str or bytes or None)
@@ -160,18 +160,19 @@
 +    **kwargs,
 +}
  a = (1,)
- b = 1,
+-b = 1,
++b = (1,)
  c = 1
  d = (1,) + a + (2,)
  e = (1,).count(1)
 -what_is_up_with_those_new_coord_names = (coord_names + set(vars_to_create)) + set(vars_to_remove)
 -what_is_up_with_those_new_coord_names = (coord_names | set(vars_to_create)) - set(vars_to_remove)
 -result = session.query(models.Customer.id).filter(models.Customer.account_id == account_id, models.Customer.email == email_address).order_by(models.Customer.id.asc(),).all()
-+what_is_up_with_those_new_coord_names = (coord_names + set(vars_to_create)) + set(
++what_is_up_with_those_new_coord_names = (
-+    vars_to_remove
++    (coord_names + set(vars_to_create)) + set(vars_to_remove)
 +)
-+what_is_up_with_those_new_coord_names = (coord_names | set(vars_to_create)) - set(
++what_is_up_with_those_new_coord_names = (
-+    vars_to_remove
++    (coord_names | set(vars_to_create)) - set(vars_to_remove)
 +)
 +result = session.query(models.Customer.id).filter(
 +    models.Customer.account_id == account_id, models.Customer.email == email_address
@@ -190,9 +191,10 @@
 +
  def gen():
      yield from outside_of_generator
-     a = (yield)
+-    a = (yield)
- 
++    a = yield
 +
+ 
  async def f():
      await some.complicated[0].call(with_args=(True or (1 is not 1)))
 -print(* [] or [1])

tests/expression.py

@@ -266,11 +266,11 @@ async def f():
 lambda a, b, c=True: a
 lambda a, b, c=True, *, d=(1 << v2), e="str": a
 lambda a, b, c=True, *vararg, d=(v1 << 2), e="str", **kwargs: a + b
-foo = (
+foo = lambda port_id, ignore_missing: {
-    lambda port_id, ignore_missing: {"port1": port1_resource, "port2": port2_resource}[
+    "port1": port1_resource, "port2": port2_resource
-        port_id
+}[
-    ]
+    port_id
-)
+]
 1 if True else 2
 str or None if True else str or bytes or None
 (str or None) if True else (str or bytes or None)
@@ -397,15 +397,15 @@ async def f():
     **kwargs,
 }
 a = (1,)
-b = 1,
+b = (1,)
 c = 1
 d = (1,) + a + (2,)
 e = (1,).count(1)
-what_is_up_with_those_new_coord_names = (coord_names + set(vars_to_create)) + set(
+what_is_up_with_those_new_coord_names = (
-    vars_to_remove
+    (coord_names + set(vars_to_create)) + set(vars_to_remove)
 )
-what_is_up_with_those_new_coord_names = (coord_names | set(vars_to_create)) - set(
+what_is_up_with_those_new_coord_names = (
-    vars_to_remove
+    (coord_names | set(vars_to_create)) - set(vars_to_remove)
 )
 result = session.query(models.Customer.id).filter(
     models.Customer.account_id == account_id, models.Customer.email == email_address
@@ -424,7 +424,7 @@ async def f():
 
 def gen():
     yield from outside_of_generator
-    a = (yield)
+    a = yield
 
 
 async def f():