"""Builds on top of nodes.py to track brackets."""
-from dataclasses import dataclass, field
import sys
-from typing import Dict, Iterable, List, Optional, Tuple, Union
+from dataclasses import dataclass, field
+from typing import Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union
if sys.version_info < (3, 8):
from typing_extensions import Final
else:
from typing import Final
-from blib2to3.pytree import Leaf, Node
+from black.nodes import (
+ BRACKET,
+ CLOSING_BRACKETS,
+ COMPARATORS,
+ LOGIC_OPERATORS,
+ MATH_OPERATORS,
+ OPENING_BRACKETS,
+ UNPACKING_PARENTS,
+ VARARGS_PARENTS,
+ is_vararg,
+ syms,
+)
from blib2to3.pgen2 import token
-
-from black.nodes import syms, is_vararg, VARARGS_PARENTS, UNPACKING_PARENTS
-from black.nodes import BRACKET, OPENING_BRACKETS, CLOSING_BRACKETS
-from black.nodes import MATH_OPERATORS, COMPARATORS, LOGIC_OPERATORS
+from blib2to3.pytree import Leaf, Node
# types
LN = Union[Leaf, Node]
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:
+ # Only track brackets from the first opening bracket to the last closing
+ # bracket.
+ 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()
+ ids = set()
+ depth = 0
+ for i in range(end_index, start_index - 1, -1):
+ 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
+ return ids