src/black/brackets.py

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