Fix a crash on dicts with paren-wrapped long string keys (#3262)

Fix a crash when formatting some dicts with parenthesis-wrapped long
string keys. When LL[0] is an atom string, we need to check the atom
node's siblings instead of LL[0] itself, e.g.:

    dictsetmaker
      atom
        STRING '"This is a really long string that can\'t be expected to fit in one line and is used as a nested dict\'s key"'
      /atom
      COLON ':'
      atom
        LSQB ' ' '['
        listmaker
          STRING '"value"'
          COMMA ','
          STRING ' ' '"value"'
        /listmaker
        RSQB ']'
      /atom
      COMMA ','
    /dictsetmaker

Author: yilei

CHANGES.md

@@ -14,6 +14,9 @@
 
 <!-- Changes that affect Black's preview style -->
 
+- Fix a crash when formatting some dicts with parenthesis-wrapped long string keys
+  (#3262)
+
 ### Configuration
 
 <!-- Changes to how Black can be configured -->

src/black/trans.py

@@ -1071,6 +1071,16 @@ def _prefer_paren_wrap_match(LL: List[Leaf]) -> Optional[int]:
             # And the string is surrounded by commas (or is the first/last child)...
             prev_sibling = LL[0].prev_sibling
             next_sibling = LL[0].next_sibling
+            if (
+                not prev_sibling
+                and not next_sibling
+                and parent_type(LL[0]) == syms.atom
+            ):
+                # If it's an atom string, we need to check the parent atom's siblings.
+                parent = LL[0].parent
+                assert parent is not None  # For type checkers.
+                prev_sibling = parent.prev_sibling
+                next_sibling = parent.next_sibling
             if (not prev_sibling or prev_sibling.type == token.COMMA) and (
                 not next_sibling or next_sibling.type == token.COMMA
             ):

tests/data/preview/long_strings.py

@@ -18,6 +18,14 @@
 
 D4 = {"A long and ridiculous {}".format(string_key): "This is a really really really long string that has to go i,side of a dictionary. It is soooo bad.", some_func("calling", "some", "stuff"): "This is a really really really long string that has to go inside of a dictionary. It is {soooo} bad (#{x}).".format(sooo="soooo", x=2), "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)}
 
+D5 = {  # Test for https://github.com/psf/black/issues/3261
+    ("This is a really long string that can't be expected to fit in one line and is used as a nested dict's key"): {"inner": "value"},
+}
+
+D6 = {  # Test for https://github.com/psf/black/issues/3261
+    ("This is a really long string that can't be expected to fit in one line and is used as a dict's key"): ["value1", "value2"],
+}
+
 L1 = ["The is a short string", "This is a really long string that can't possibly be expected to fit all together on one line. Also it is inside a list literal, so it's expected to be wrapped in parens when spliting to avoid implicit str concatenation.", short_call("arg", {"key": "value"}), "This is another really really (not really) long string that also can't be expected to fit on one line and is, like the other string, inside a list literal.", ("parens should be stripped for short string in list")]
 
 L2 = ["This is a really long string that can't be expected to fit in one line and is the only child of a list literal."]
@@ -357,6 +365,19 @@ def foo():
     % ("soooo", 2),
 }
 
+D5 = {  # Test for https://github.com/psf/black/issues/3261
+    "This is a really long string that can't be expected to fit in one line and is used as a nested dict's key": {
+        "inner": "value"
+    },
+}
+
+D6 = {  # Test for https://github.com/psf/black/issues/3261
+    "This is a really long string that can't be expected to fit in one line and is used as a dict's key": [
+        "value1",
+        "value2",
+    ],
+}
+
 L1 = [
     "The is a short string",
     (