X-Git-Url: https://git.madduck.net/etc/vim.git/blobdiff_plain/117891878e5be4d6b771ae5de299e51b679cea27..d4a85643a465f5fae2113d07d22d021d4af4795a:/src/black/brackets.py?ds=sidebyside diff --git a/src/black/brackets.py b/src/black/brackets.py index c5ed4bf..0a5317f 100644 --- a/src/black/brackets.py +++ b/src/black/brackets.py @@ -1,20 +1,28 @@ """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] @@ -332,3 +340,35 @@ 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: + # 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