src/black/brackets.py

   1 """Builds on top of nodes.py to track brackets."""
   2
   3 import sys
   4 from dataclasses import dataclass, field
   5 from typing import Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union
   6
   7 if sys.version_info < (3, 8):
   8     from typing_extensions import Final
   9 else:
  10     from typing import Final
  11
  12 from black.nodes import (
  13     BRACKET,
  14     CLOSING_BRACKETS,
  15     COMPARATORS,
  16     LOGIC_OPERATORS,
  17     MATH_OPERATORS,
  18     OPENING_BRACKETS,
  19     UNPACKING_PARENTS,
  20     VARARGS_PARENTS,
  21     is_vararg,
  22     syms,
  23 )
  24 from blib2to3.pgen2 import token
  25 from blib2to3.pytree import Leaf, Node
  26
  27 # types
  28 LN = Union[Leaf, Node]
  29 Depth = int
  30 LeafID = int
  31 NodeType = int
  32 Priority = int
  33
  34
  35 COMPREHENSION_PRIORITY: Final = 20
  36 COMMA_PRIORITY: Final = 18
  37 TERNARY_PRIORITY: Final = 16
  38 LOGIC_PRIORITY: Final = 14
  39 STRING_PRIORITY: Final = 12
  40 COMPARATOR_PRIORITY: Final = 10
  41 MATH_PRIORITIES: Final = {
  42     token.VBAR: 9,
  43     token.CIRCUMFLEX: 8,
  44     token.AMPER: 7,
  45     token.LEFTSHIFT: 6,
  46     token.RIGHTSHIFT: 6,
  47     token.PLUS: 5,
  48     token.MINUS: 5,
  49     token.STAR: 4,
  50     token.SLASH: 4,
  51     token.DOUBLESLASH: 4,
  52     token.PERCENT: 4,
  53     token.AT: 4,
  54     token.TILDE: 3,
  55     token.DOUBLESTAR: 2,
  56 }
  57 DOT_PRIORITY: Final = 1
  58
  59
  60 class BracketMatchError(Exception):
  61     """Raised when an opening bracket is unable to be matched to a closing bracket."""
  62
  63
  64 @dataclass
  65 class BracketTracker:
  66     """Keeps track of brackets on a line."""
  67
  68     depth: int = 0
  69     bracket_match: Dict[Tuple[Depth, NodeType], Leaf] = field(default_factory=dict)
  70     delimiters: Dict[LeafID, Priority] = field(default_factory=dict)
  71     previous: Optional[Leaf] = None
  72     _for_loop_depths: List[int] = field(default_factory=list)
  73     _lambda_argument_depths: List[int] = field(default_factory=list)
  74     invisible: List[Leaf] = field(default_factory=list)
  75
  76     def mark(self, leaf: Leaf) -> None:
  77         """Mark `leaf` with bracket-related metadata. Keep track of delimiters.
  78
  79         All leaves receive an int `bracket_depth` field that stores how deep
  80         within brackets a given leaf is. 0 means there are no enclosing brackets
  81         that started on this line.
  82
  83         If a leaf is itself a closing bracket, it receives an `opening_bracket`
  84         field that it forms a pair with. This is a one-directional link to
  85         avoid reference cycles.
  86
  87         If a leaf is a delimiter (a token on which Black can split the line if
  88         needed) and it's on depth 0, its `id()` is stored in the tracker's
  89         `delimiters` field.
  90         """
  91         if leaf.type == token.COMMENT:
  92             return
  93
  94         self.maybe_decrement_after_for_loop_variable(leaf)
  95         self.maybe_decrement_after_lambda_arguments(leaf)
  96         if leaf.type in CLOSING_BRACKETS:
  97             self.depth -= 1
  98             try:
  99                 opening_bracket = self.bracket_match.pop((self.depth, leaf.type))
 100             except KeyError as e:
 101                 raise BracketMatchError(
 102                     "Unable to match a closing bracket to the following opening"
 103                     f" bracket: {leaf}"
 104                 ) from e
 105             leaf.opening_bracket = opening_bracket
 106             if not leaf.value:
 107                 self.invisible.append(leaf)
 108         leaf.bracket_depth = self.depth
 109         if self.depth == 0:
 110             delim = is_split_before_delimiter(leaf, self.previous)
 111             if delim and self.previous is not None:
 112                 self.delimiters[id(self.previous)] = delim
 113             else:
 114                 delim = is_split_after_delimiter(leaf, self.previous)
 115                 if delim:
 116                     self.delimiters[id(leaf)] = delim
 117         if leaf.type in OPENING_BRACKETS:
 118             self.bracket_match[self.depth, BRACKET[leaf.type]] = leaf
 119             self.depth += 1
 120             if not leaf.value:
 121                 self.invisible.append(leaf)
 122         self.previous = leaf
 123         self.maybe_increment_lambda_arguments(leaf)
 124         self.maybe_increment_for_loop_variable(leaf)
 125
 126     def any_open_brackets(self) -> bool:
 127         """Return True if there is an yet unmatched open bracket on the line."""
 128         return bool(self.bracket_match)
 129
 130     def max_delimiter_priority(self, exclude: Iterable[LeafID] = ()) -> Priority:
 131         """Return the highest priority of a delimiter found on the line.
 132
 133         Values are consistent with what `is_split_*_delimiter()` return.
 134         Raises ValueError on no delimiters.
 135         """
 136         return max(v for k, v in self.delimiters.items() if k not in exclude)
 137
 138     def delimiter_count_with_priority(self, priority: Priority = 0) -> int:
 139         """Return the number of delimiters with the given `priority`.
 140
 141         If no `priority` is passed, defaults to max priority on the line.
 142         """
 143         if not self.delimiters:
 144             return 0
 145
 146         priority = priority or self.max_delimiter_priority()
 147         return sum(1 for p in self.delimiters.values() if p == priority)
 148
 149     def maybe_increment_for_loop_variable(self, leaf: Leaf) -> bool:
 150         """In a for loop, or comprehension, the variables are often unpacks.
 151
 152         To avoid splitting on the comma in this situation, increase the depth of
 153         tokens between `for` and `in`.
 154         """
 155         if leaf.type == token.NAME and leaf.value == "for":
 156             self.depth += 1
 157             self._for_loop_depths.append(self.depth)
 158             return True
 159
 160         return False
 161
 162     def maybe_decrement_after_for_loop_variable(self, leaf: Leaf) -> bool:
 163         """See `maybe_increment_for_loop_variable` above for explanation."""
 164         if (
 165             self._for_loop_depths
 166             and self._for_loop_depths[-1] == self.depth
 167             and leaf.type == token.NAME
 168             and leaf.value == "in"
 169         ):
 170             self.depth -= 1
 171             self._for_loop_depths.pop()
 172             return True
 173
 174         return False
 175
 176     def maybe_increment_lambda_arguments(self, leaf: Leaf) -> bool:
 177         """In a lambda expression, there might be more than one argument.
 178
 179         To avoid splitting on the comma in this situation, increase the depth of
 180         tokens between `lambda` and `:`.
 181         """
 182         if leaf.type == token.NAME and leaf.value == "lambda":
 183             self.depth += 1
 184             self._lambda_argument_depths.append(self.depth)
 185             return True
 186
 187         return False
 188
 189     def maybe_decrement_after_lambda_arguments(self, leaf: Leaf) -> bool:
 190         """See `maybe_increment_lambda_arguments` above for explanation."""
 191         if (
 192             self._lambda_argument_depths
 193             and self._lambda_argument_depths[-1] == self.depth
 194             and leaf.type == token.COLON
 195         ):
 196             self.depth -= 1
 197             self._lambda_argument_depths.pop()
 198             return True
 199
 200         return False
 201
 202     def get_open_lsqb(self) -> Optional[Leaf]:
 203         """Return the most recent opening square bracket (if any)."""
 204         return self.bracket_match.get((self.depth - 1, token.RSQB))
 205
 206
 207 def is_split_after_delimiter(leaf: Leaf, previous: Optional[Leaf] = None) -> Priority:
 208     """Return the priority of the `leaf` delimiter, given a line break after it.
 209
 210     The delimiter priorities returned here are from those delimiters that would
 211     cause a line break after themselves.
 212
 213     Higher numbers are higher priority.
 214     """
 215     if leaf.type == token.COMMA:
 216         return COMMA_PRIORITY
 217
 218     return 0
 219
 220
 221 def is_split_before_delimiter(leaf: Leaf, previous: Optional[Leaf] = None) -> Priority:
 222     """Return the priority of the `leaf` delimiter, given a line break before it.
 223
 224     The delimiter priorities returned here are from those delimiters that would
 225     cause a line break before themselves.
 226
 227     Higher numbers are higher priority.
 228     """
 229     if is_vararg(leaf, within=VARARGS_PARENTS | UNPACKING_PARENTS):
 230         # * and ** might also be MATH_OPERATORS but in this case they are not.
 231         # Don't treat them as a delimiter.
 232         return 0
 233
 234     if (
 235         leaf.type == token.DOT
 236         and leaf.parent
 237         and leaf.parent.type not in {syms.import_from, syms.dotted_name}
 238         and (previous is None or previous.type in CLOSING_BRACKETS)
 239     ):
 240         return DOT_PRIORITY
 241
 242     if (
 243         leaf.type in MATH_OPERATORS
 244         and leaf.parent
 245         and leaf.parent.type not in {syms.factor, syms.star_expr}
 246     ):
 247         return MATH_PRIORITIES[leaf.type]
 248
 249     if leaf.type in COMPARATORS:
 250         return COMPARATOR_PRIORITY
 251
 252     if (
 253         leaf.type == token.STRING
 254         and previous is not None
 255         and previous.type == token.STRING
 256     ):
 257         return STRING_PRIORITY
 258
 259     if leaf.type not in {token.NAME, token.ASYNC}:
 260         return 0
 261
 262     if (
 263         leaf.value == "for"
 264         and leaf.parent
 265         and leaf.parent.type in {syms.comp_for, syms.old_comp_for}
 266         or leaf.type == token.ASYNC
 267     ):
 268         if (
 269             not isinstance(leaf.prev_sibling, Leaf)
 270             or leaf.prev_sibling.value != "async"
 271         ):
 272             return COMPREHENSION_PRIORITY
 273
 274     if (
 275         leaf.value == "if"
 276         and leaf.parent
 277         and leaf.parent.type in {syms.comp_if, syms.old_comp_if}
 278     ):
 279         return COMPREHENSION_PRIORITY
 280
 281     if leaf.value in {"if", "else"} and leaf.parent and leaf.parent.type == syms.test:
 282         return TERNARY_PRIORITY
 283
 284     if leaf.value == "is":
 285         return COMPARATOR_PRIORITY
 286
 287     if (
 288         leaf.value == "in"
 289         and leaf.parent
 290         and leaf.parent.type in {syms.comp_op, syms.comparison}
 291         and not (
 292             previous is not None
 293             and previous.type == token.NAME
 294             and previous.value == "not"
 295         )
 296     ):
 297         return COMPARATOR_PRIORITY
 298
 299     if (
 300         leaf.value == "not"
 301         and leaf.parent
 302         and leaf.parent.type == syms.comp_op
 303         and not (
 304             previous is not None
 305             and previous.type == token.NAME
 306             and previous.value == "is"
 307         )
 308     ):
 309         return COMPARATOR_PRIORITY
 310
 311     if leaf.value in LOGIC_OPERATORS and leaf.parent:
 312         return LOGIC_PRIORITY
 313
 314     return 0
 315
 316
 317 def max_delimiter_priority_in_atom(node: LN) -> Priority:
 318     """Return maximum delimiter priority inside `node`.
 319
 320     This is specific to atoms with contents contained in a pair of parentheses.
 321     If `node` isn't an atom or there are no enclosing parentheses, returns 0.
 322     """
 323     if node.type != syms.atom:
 324         return 0
 325
 326     first = node.children[0]
 327     last = node.children[-1]
 328     if not (first.type == token.LPAR and last.type == token.RPAR):
 329         return 0
 330
 331     bt = BracketTracker()
 332     for c in node.children[1:-1]:
 333         if isinstance(c, Leaf):
 334             bt.mark(c)
 335         else:
 336             for leaf in c.leaves():
 337                 bt.mark(leaf)
 338     try:
 339         return bt.max_delimiter_priority()
 340
 341     except ValueError:
 342         return 0
 343
 344
 345 def get_leaves_inside_matching_brackets(leaves: Sequence[Leaf]) -> Set[LeafID]:
 346     """Return leaves that are inside matching brackets.
 347
 348     The input `leaves` can have non-matching brackets at the head or tail parts.
 349     Matching brackets are included.
 350     """
 351     try:
 352         # Start with the first opening bracket and ignore closing brackets before.
 353         start_index = next(
 354             i for i, l in enumerate(leaves) if l.type in OPENING_BRACKETS
 355         )
 356     except StopIteration:
 357         return set()
 358     bracket_stack = []
 359     ids = set()
 360     for i in range(start_index, len(leaves)):
 361         leaf = leaves[i]
 362         if leaf.type in OPENING_BRACKETS:
 363             bracket_stack.append((BRACKET[leaf.type], i))
 364         if leaf.type in CLOSING_BRACKETS:
 365             if bracket_stack and leaf.type == bracket_stack[-1][0]:
 366                 _, start = bracket_stack.pop()
 367                 for j in range(start, i + 1):
 368                     ids.add(id(leaves[j]))
 369             else:
 370                 break
 371     return ids