From 96e62c57e3023977de177a8ba34678007a63f1fe Mon Sep 17 00:00:00 2001
From: "Yilei \"Dolee\" Yang" <yileiyang@google.com>
Date: Sat, 10 Dec 2022 07:58:45 -0800
Subject: [PATCH 1/1] Fix a crash in preview style with assert + parenthesized
 string. (#3415)

The bug is in the `get_leaves_inside_matching_brackets` on the third line below:

```python
assert xxxxxxxxx.xxxxxxxxx.xxxxxxxxx(
    xxxxxxxxx
).xxxxxxxxxxxxxxxxxx(), (
    "xxx {xxxxxxxxx} xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
)
```

Including the invisible paren, third line is `).xxxxxxxxxxxxxxxxxx()), (`, that it has a matched pair then an unmatched closing paren afterwards. This PR ensures the returned leaves are actually matched.

Fixes #3414.
---
 CHANGES.md                                    |  2 ++
 src/black/brackets.py                         | 25 ++++++++-----------
 .../trailing_commas_in_leading_parts.py       | 14 +++++++++++
 3 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 20ae65d..c84feb0 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -14,6 +14,8 @@
 
 <!-- Changes that affect Black's preview style -->
 
+- Fix a crash in preview style with assert + parenthesized string (#3415)
+
 ### Configuration
 
 <!-- Changes to how Black can be configured -->
diff --git a/src/black/brackets.py b/src/black/brackets.py
index 0a5317f..ec9708c 100644
--- a/src/black/brackets.py
+++ b/src/black/brackets.py
@@ -349,26 +349,23 @@ def get_leaves_inside_matching_brackets(leaves: Sequence[Leaf]) -> Set[LeafID]:
     Matching brackets are included.
     """
     try:
-        # Only track brackets from the first opening bracket to the last closing
-        # bracket.
+        # Start with the first opening bracket and ignore closing brackets before.
         start_index = next(
             i for i, l in enumerate(leaves) if l.type in OPENING_BRACKETS
         )
-        end_index = next(
-            len(leaves) - i
-            for i, l in enumerate(reversed(leaves))
-            if l.type in CLOSING_BRACKETS
-        )
     except StopIteration:
         return set()
+    bracket_stack = []
     ids = set()
-    depth = 0
-    for i in range(end_index, start_index - 1, -1):
+    for i in range(start_index, len(leaves)):
         leaf = leaves[i]
-        if leaf.type in CLOSING_BRACKETS:
-            depth += 1
-        if depth > 0:
-            ids.add(id(leaf))
         if leaf.type in OPENING_BRACKETS:
-            depth -= 1
+            bracket_stack.append((BRACKET[leaf.type], i))
+        if leaf.type in CLOSING_BRACKETS:
+            if bracket_stack and leaf.type == bracket_stack[-1][0]:
+                _, start = bracket_stack.pop()
+                for j in range(start, i + 1):
+                    ids.add(id(leaves[j]))
+            else:
+                break
     return ids
diff --git a/tests/data/preview/trailing_commas_in_leading_parts.py b/tests/data/preview/trailing_commas_in_leading_parts.py
index 676725c..99d82a6 100644
--- a/tests/data/preview/trailing_commas_in_leading_parts.py
+++ b/tests/data/preview/trailing_commas_in_leading_parts.py
@@ -25,6 +25,13 @@ assert (
     == long_module.long_class.long_func()["some_key"].another_func(arg1)
 )
 
+# Regression test for https://github.com/psf/black/issues/3414.
+assert xxxxxxxxx.xxxxxxxxx.xxxxxxxxx(
+    xxxxxxxxx
+).xxxxxxxxxxxxxxxxxx(), (
+    "xxx {xxxxxxxxx} xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
+)
+
 
 # output
 
@@ -72,3 +79,10 @@ assert (
     long_module.long_class.long_func().another_func()
     == long_module.long_class.long_func()["some_key"].another_func(arg1)
 )
+
+# Regression test for https://github.com/psf/black/issues/3414.
+assert xxxxxxxxx.xxxxxxxxx.xxxxxxxxx(
+    xxxxxxxxx
+).xxxxxxxxxxxxxxxxxx(), (
+    "xxx {xxxxxxxxx} xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
+)
-- 
2.39.5