Fix a crash in preview style with assert + parenthesized string. (#3415)

[etc/vim.git] / src / black / brackets.py

diff --git a/src/black/brackets.py b/src/black/brackets.py
index 3566f5b6c3757f3f12459e4c2f379b06da9e062c..ec9708cb08a214d55b8680ac07dd0ebf0fde413a 100644 (file)
--- a/src/black/brackets.py
+++ b/src/black/brackets.py
@@ -2,7 +2,7 @@
 
 import sys
 from dataclasses import dataclass, field
-from typing import Dict, Iterable, List, Optional, Tuple, Union
+from typing import Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union
 
 if sys.version_info < (3, 8):
     from typing_extensions import Final
@@ -340,3 +340,32 @@ def max_delimiter_priority_in_atom(node: LN) -> Priority:
 
     except ValueError:
         return 0
+
+
+def get_leaves_inside_matching_brackets(leaves: Sequence[Leaf]) -> Set[LeafID]:
+    """Return leaves that are inside matching brackets.
+
+    The input `leaves` can have non-matching brackets at the head or tail parts.
+    Matching brackets are included.
+    """
+    try:
+        # 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
+        )
+    except StopIteration:
+        return set()
+    bracket_stack = []
+    ids = set()
+    for i in range(start_index, len(leaves)):
+        leaf = leaves[i]
+        if leaf.type in OPENING_BRACKETS:
+            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