Improve long values in dict literals (#3440)

CHANGES.md

@@ -19,6 +19,9 @@
 - Fix a crash in preview style with assert + parenthesized string (#3415)
 - Do not put the closing quotes in a docstring on a separate line, even if the line is
   too long (#3430)
+- Long values in dict literals are now wrapped in parentheses; correspondingly
+  unnecessary parentheses around short values in dict literals are now removed; long
+  string lambda values are now wrapped in parentheses (#3440)
 
 ### Configuration
 

src/black/linegen.py

@@ -179,6 +179,23 @@ def visit_stmt(
 
             yield from self.visit(child)
 
+    def visit_dictsetmaker(self, node: Node) -> Iterator[Line]:
+        if Preview.wrap_long_dict_values_in_parens in self.mode:
+            for i, child in enumerate(node.children):
+                if i == 0:
+                    continue
+                if node.children[i - 1].type == token.COLON:
+                    if child.type == syms.atom and child.children[0].type == token.LPAR:
+                        if maybe_make_parens_invisible_in_atom(
+                            child,
+                            parent=node,
+                            remove_brackets_around_comma=False,
+                        ):
+                            wrap_in_parentheses(node, child, visible=False)
+                    else:
+                        wrap_in_parentheses(node, child, visible=False)
+        yield from self.visit_default(node)
+
     def visit_funcdef(self, node: Node) -> Iterator[Line]:
         """Visit function definition."""
         if Preview.annotation_parens not in self.mode:

src/black/mode.py

@@ -157,8 +157,11 @@ class Preview(Enum):
     one_element_subscript = auto()
     remove_block_trailing_newline = auto()
     remove_redundant_parens = auto()
+    # NOTE: string_processing requires wrap_long_dict_values_in_parens
+    # for https://github.com/psf/black/issues/3117 to be fixed.
     string_processing = auto()
     skip_magic_trailing_comma_in_subscript = auto()
+    wrap_long_dict_values_in_parens = auto()
 
 
 class Deprecated(UserWarning):

src/black/trans.py

@@ -1638,6 +1638,8 @@ class StringParenWrapper(BaseStringSplitter, CustomSplitMapMixin):
         * The line is a dictionary key assignment where some valid key is being
         assigned the value of some string.
             OR
+        * The line is an lambda expression and the value is a string.
+            OR
         * The line starts with an "atom" string that prefers to be wrapped in
         parens. It's preferred to be wrapped when the string is surrounded by
         commas (or is the first/last child).
@@ -1683,7 +1685,7 @@ def do_splitter_match(self, line: Line) -> TMatchResult:
             or self._else_match(LL)
             or self._assert_match(LL)
             or self._assign_match(LL)
-            or self._dict_match(LL)
+            or self._dict_or_lambda_match(LL)
             or self._prefer_paren_wrap_match(LL)
         )
 
@@ -1841,22 +1843,23 @@ def _assign_match(LL: List[Leaf]) -> Optional[int]:
         return None
 
     @staticmethod
-    def _dict_match(LL: List[Leaf]) -> Optional[int]:
+    def _dict_or_lambda_match(LL: List[Leaf]) -> Optional[int]:
         """
         Returns:
             string_idx such that @LL[string_idx] is equal to our target (i.e.
             matched) string, if this line matches the dictionary key assignment
-            statement requirements listed in the 'Requirements' section of this
-            classes' docstring.
+            statement or lambda expression requirements listed in the
+            'Requirements' section of this classes' docstring.
                 OR
             None, otherwise.
         """
-        # If this line is apart of a dictionary key assignment...
-        if syms.dictsetmaker in [parent_type(LL[0]), parent_type(LL[0].parent)]:
+        # If this line is a part of a dictionary key assignment or lambda expression...
+        parent_types = [parent_type(LL[0]), parent_type(LL[0].parent)]
+        if syms.dictsetmaker in parent_types or syms.lambdef in parent_types:
             is_valid_index = is_valid_index_factory(LL)
 
             for i, leaf in enumerate(LL):
-                # We MUST find a colon...
+                # We MUST find a colon, it can either be dict's or lambda's colon...
                 if leaf.type == token.COLON:
                     idx = i + 2 if is_empty_par(LL[i + 1]) else i + 1
 
@@ -1951,6 +1954,25 @@ def do_transform(self, line: Line, string_idx: int) -> Iterator[TResult[Line]]:
                     f" (left_leaves={left_leaves}, right_leaves={right_leaves})"
                 )
                 old_rpar_leaf = right_leaves.pop()
+            elif right_leaves and right_leaves[-1].type == token.RPAR:
+                # Special case for lambda expressions as dict's value, e.g.:
+                #     my_dict = {
+                #        "key": lambda x: f"formatted: {x},
+                #     }
+                # After wrapping the dict's value with parentheses, the string is
+                # followed by a RPAR but its opening bracket is lambda's, not
+                # the string's:
+                #        "key": (lambda x: f"formatted: {x}),
+                opening_bracket = right_leaves[-1].opening_bracket
+                if opening_bracket is not None and opening_bracket in left_leaves:
+                    index = left_leaves.index(opening_bracket)
+                    if (
+                        index > 0
+                        and index < len(left_leaves) - 1
+                        and left_leaves[index - 1].type == token.COLON
+                        and left_leaves[index + 1].value == "lambda"
+                    ):
+                        right_leaves.pop()
 
             append_leaves(string_line, line, right_leaves)
 

tests/data/preview/long_dict_values.py

@@ -0,0 +1,53 @@
+my_dict = {
+    "something_something":
+        r"Lorem ipsum dolor sit amet, an sed convenire eloquentiam \t"
+        r"signiferumque, duo ea vocibus consetetur scriptorem. Facer \t"
+        r"signiferumque, duo ea vocibus consetetur scriptorem. Facer \t",
+}
+
+my_dict = {
+    "a key in my dict": a_very_long_variable * and_a_very_long_function_call() / 100000.0
+}
+
+my_dict = {
+    "a key in my dict": a_very_long_variable * and_a_very_long_function_call() * and_another_long_func() / 100000.0
+}
+
+my_dict = {
+    "a key in my dict": MyClass.some_attribute.first_call().second_call().third_call(some_args="some value")
+}
+
+
+# output
+
+
+my_dict = {
+    "something_something": (
+        r"Lorem ipsum dolor sit amet, an sed convenire eloquentiam \t"
+        r"signiferumque, duo ea vocibus consetetur scriptorem. Facer \t"
+        r"signiferumque, duo ea vocibus consetetur scriptorem. Facer \t"
+    ),
+}
+
+my_dict = {
+    "a key in my dict": (
+        a_very_long_variable * and_a_very_long_function_call() / 100000.0
+    )
+}
+
+my_dict = {
+    "a key in my dict": (
+        a_very_long_variable
+        * and_a_very_long_function_call()
+        * and_another_long_func()
+        / 100000.0
+    )
+}
+
+my_dict = {
+    "a key in my dict": (
+        MyClass.some_attribute.first_call()
+        .second_call()
+        .third_call(some_args="some value")
+    )
+}

tests/data/preview/long_strings.py

@@ -278,6 +278,15 @@ def foo():
     "........................................................................... \\N{LAO KO LA}"
 )
 
+msg = lambda x: f"this is a very very very long lambda value {x} that doesn't fit on a single line"
+
+dict_with_lambda_values = {
+    "join": lambda j: (
+        f"{j.__class__.__name__}({some_function_call(j.left)}, "
+        f"{some_function_call(j.right)})"
+    ),
+}
+
 
 # output
 
@@ -362,9 +371,8 @@ def foo():
     "A %s %s"
     % ("formatted", "string"): (
         "This is a really really really long string that has to go inside of a"
-        " dictionary. It is %s bad (#%d)."
-    )
-    % ("soooo", 2),
+        " dictionary. It is %s bad (#%d)." % ("soooo", 2)
+    ),
 }
 
 D5 = {  # Test for https://github.com/psf/black/issues/3261
@@ -806,3 +814,17 @@ def foo():
     "..........................................................................."
     " \\N{LAO KO LA}"
 )
+
+msg = (
+    lambda x: (
+        f"this is a very very very long lambda value {x} that doesn't fit on a single"
+        " line"
+    )
+)
+
+dict_with_lambda_values = {
+    "join": lambda j: (
+        f"{j.__class__.__name__}({some_function_call(j.left)}, "
+        f"{some_function_call(j.right)})"
+    ),
+}

tests/data/preview/long_strings__regression.py

@@ -524,6 +524,13 @@ async def foo(self):
     },
 )
 
+# Regression test for https://github.com/psf/black/issues/3117.
+some_dict = {
+    "something_something":
+        r"Lorem ipsum dolor sit amet, an sed convenire eloquentiam \t"
+        r"signiferumque, duo ea vocibus consetetur scriptorem. Facer \t",
+}
+
 
 # output
 
@@ -1178,3 +1185,11 @@ async def foo(self):
         ),
     },
 )
+
+# Regression test for https://github.com/psf/black/issues/3117.
+some_dict = {
+    "something_something": (
+        r"Lorem ipsum dolor sit amet, an sed convenire eloquentiam \t"
+        r"signiferumque, duo ea vocibus consetetur scriptorem. Facer \t"
+    ),
+}