1 """Builds on top of nodes.py to track brackets."""
3 from dataclasses
import dataclass
, field
4 from typing
import Dict
, Final
, Iterable
, List
, Optional
, Sequence
, Set
, Tuple
, Union
6 from black
.nodes
import (
18 from blib2to3
.pgen2
import token
19 from blib2to3
.pytree
import Leaf
, Node
22 LN
= Union
[Leaf
, Node
]
29 COMPREHENSION_PRIORITY
: Final
= 20
30 COMMA_PRIORITY
: Final
= 18
31 TERNARY_PRIORITY
: Final
= 16
32 LOGIC_PRIORITY
: Final
= 14
33 STRING_PRIORITY
: Final
= 12
34 COMPARATOR_PRIORITY
: Final
= 10
35 MATH_PRIORITIES
: Final
= {
51 DOT_PRIORITY
: Final
= 1
54 class BracketMatchError(Exception):
55 """Raised when an opening bracket is unable to be matched to a closing bracket."""
60 """Keeps track of brackets on a line."""
63 bracket_match
: Dict
[Tuple
[Depth
, NodeType
], Leaf
] = field(default_factory
=dict)
64 delimiters
: Dict
[LeafID
, Priority
] = field(default_factory
=dict)
65 previous
: Optional
[Leaf
] = None
66 _for_loop_depths
: List
[int] = field(default_factory
=list)
67 _lambda_argument_depths
: List
[int] = field(default_factory
=list)
68 invisible
: List
[Leaf
] = field(default_factory
=list)
70 def mark(self
, leaf
: Leaf
) -> None:
71 """Mark `leaf` with bracket-related metadata. Keep track of delimiters.
73 All leaves receive an int `bracket_depth` field that stores how deep
74 within brackets a given leaf is. 0 means there are no enclosing brackets
75 that started on this line.
77 If a leaf is itself a closing bracket and there is a matching opening
78 bracket earlier, it receives an `opening_bracket` field with which it forms a
79 pair. This is a one-directional link to avoid reference cycles. Closing
80 bracket without opening happens on lines continued from previous
81 breaks, e.g. `) -> "ReturnType":` as part of a funcdef where we place
82 the return type annotation on its own line of the previous closing RPAR.
84 If a leaf is a delimiter (a token on which Black can split the line if
85 needed) and it's on depth 0, its `id()` is stored in the tracker's
88 if leaf
.type == token
.COMMENT
:
93 and leaf
.type in CLOSING_BRACKETS
94 and (self
.depth
, leaf
.type) not in self
.bracket_match
98 self
.maybe_decrement_after_for_loop_variable(leaf
)
99 self
.maybe_decrement_after_lambda_arguments(leaf
)
100 if leaf
.type in CLOSING_BRACKETS
:
103 opening_bracket
= self
.bracket_match
.pop((self
.depth
, leaf
.type))
104 except KeyError as e
:
105 raise BracketMatchError(
106 "Unable to match a closing bracket to the following opening"
109 leaf
.opening_bracket
= opening_bracket
111 self
.invisible
.append(leaf
)
112 leaf
.bracket_depth
= self
.depth
114 delim
= is_split_before_delimiter(leaf
, self
.previous
)
115 if delim
and self
.previous
is not None:
116 self
.delimiters
[id(self
.previous
)] = delim
118 delim
= is_split_after_delimiter(leaf
, self
.previous
)
120 self
.delimiters
[id(leaf
)] = delim
121 if leaf
.type in OPENING_BRACKETS
:
122 self
.bracket_match
[self
.depth
, BRACKET
[leaf
.type]] = leaf
125 self
.invisible
.append(leaf
)
127 self
.maybe_increment_lambda_arguments(leaf
)
128 self
.maybe_increment_for_loop_variable(leaf
)
130 def any_open_brackets(self
) -> bool:
131 """Return True if there is an yet unmatched open bracket on the line."""
132 return bool(self
.bracket_match
)
134 def max_delimiter_priority(self
, exclude
: Iterable
[LeafID
] = ()) -> Priority
:
135 """Return the highest priority of a delimiter found on the line.
137 Values are consistent with what `is_split_*_delimiter()` return.
138 Raises ValueError on no delimiters.
140 return max(v
for k
, v
in self
.delimiters
.items() if k
not in exclude
)
142 def delimiter_count_with_priority(self
, priority
: Priority
= 0) -> int:
143 """Return the number of delimiters with the given `priority`.
145 If no `priority` is passed, defaults to max priority on the line.
147 if not self
.delimiters
:
150 priority
= priority
or self
.max_delimiter_priority()
151 return sum(1 for p
in self
.delimiters
.values() if p
== priority
)
153 def maybe_increment_for_loop_variable(self
, leaf
: Leaf
) -> bool:
154 """In a for loop, or comprehension, the variables are often unpacks.
156 To avoid splitting on the comma in this situation, increase the depth of
157 tokens between `for` and `in`.
159 if leaf
.type == token
.NAME
and leaf
.value
== "for":
161 self
._for
_loop
_depths
.append(self
.depth
)
166 def maybe_decrement_after_for_loop_variable(self
, leaf
: Leaf
) -> bool:
167 """See `maybe_increment_for_loop_variable` above for explanation."""
169 self
._for
_loop
_depths
170 and self
._for
_loop
_depths
[-1] == self
.depth
171 and leaf
.type == token
.NAME
172 and leaf
.value
== "in"
175 self
._for
_loop
_depths
.pop()
180 def maybe_increment_lambda_arguments(self
, leaf
: Leaf
) -> bool:
181 """In a lambda expression, there might be more than one argument.
183 To avoid splitting on the comma in this situation, increase the depth of
184 tokens between `lambda` and `:`.
186 if leaf
.type == token
.NAME
and leaf
.value
== "lambda":
188 self
._lambda
_argument
_depths
.append(self
.depth
)
193 def maybe_decrement_after_lambda_arguments(self
, leaf
: Leaf
) -> bool:
194 """See `maybe_increment_lambda_arguments` above for explanation."""
196 self
._lambda
_argument
_depths
197 and self
._lambda
_argument
_depths
[-1] == self
.depth
198 and leaf
.type == token
.COLON
201 self
._lambda
_argument
_depths
.pop()
206 def get_open_lsqb(self
) -> Optional
[Leaf
]:
207 """Return the most recent opening square bracket (if any)."""
208 return self
.bracket_match
.get((self
.depth
- 1, token
.RSQB
))
211 def is_split_after_delimiter(leaf
: Leaf
, previous
: Optional
[Leaf
] = None) -> Priority
:
212 """Return the priority of the `leaf` delimiter, given a line break after it.
214 The delimiter priorities returned here are from those delimiters that would
215 cause a line break after themselves.
217 Higher numbers are higher priority.
219 if leaf
.type == token
.COMMA
:
220 return COMMA_PRIORITY
225 def is_split_before_delimiter(leaf
: Leaf
, previous
: Optional
[Leaf
] = None) -> Priority
:
226 """Return the priority of the `leaf` delimiter, given a line break before it.
228 The delimiter priorities returned here are from those delimiters that would
229 cause a line break before themselves.
231 Higher numbers are higher priority.
233 if is_vararg(leaf
, within
=VARARGS_PARENTS | UNPACKING_PARENTS
):
234 # * and ** might also be MATH_OPERATORS but in this case they are not.
235 # Don't treat them as a delimiter.
239 leaf
.type == token
.DOT
241 and leaf
.parent
.type not in {syms
.import_from
, syms
.dotted_name
}
242 and (previous
is None or previous
.type in CLOSING_BRACKETS
)
247 leaf
.type in MATH_OPERATORS
249 and leaf
.parent
.type not in {syms
.factor
, syms
.star_expr
}
251 return MATH_PRIORITIES
[leaf
.type]
253 if leaf
.type in COMPARATORS
:
254 return COMPARATOR_PRIORITY
257 leaf
.type == token
.STRING
258 and previous
is not None
259 and previous
.type == token
.STRING
261 return STRING_PRIORITY
263 if leaf
.type not in {token
.NAME
, token
.ASYNC
}:
269 and leaf
.parent
.type in {syms
.comp_for
, syms
.old_comp_for
}
270 or leaf
.type == token
.ASYNC
273 not isinstance(leaf
.prev_sibling
, Leaf
)
274 or leaf
.prev_sibling
.value
!= "async"
276 return COMPREHENSION_PRIORITY
281 and leaf
.parent
.type in {syms
.comp_if
, syms
.old_comp_if
}
283 return COMPREHENSION_PRIORITY
285 if leaf
.value
in {"if", "else"} and leaf
.parent
and leaf
.parent
.type == syms
.test
:
286 return TERNARY_PRIORITY
288 if leaf
.value
== "is":
289 return COMPARATOR_PRIORITY
294 and leaf
.parent
.type in {syms
.comp_op
, syms
.comparison
}
297 and previous
.type == token
.NAME
298 and previous
.value
== "not"
301 return COMPARATOR_PRIORITY
306 and leaf
.parent
.type == syms
.comp_op
309 and previous
.type == token
.NAME
310 and previous
.value
== "is"
313 return COMPARATOR_PRIORITY
315 if leaf
.value
in LOGIC_OPERATORS
and leaf
.parent
:
316 return LOGIC_PRIORITY
321 def max_delimiter_priority_in_atom(node
: LN
) -> Priority
:
322 """Return maximum delimiter priority inside `node`.
324 This is specific to atoms with contents contained in a pair of parentheses.
325 If `node` isn't an atom or there are no enclosing parentheses, returns 0.
327 if node
.type != syms
.atom
:
330 first
= node
.children
[0]
331 last
= node
.children
[-1]
332 if not (first
.type == token
.LPAR
and last
.type == token
.RPAR
):
335 bt
= BracketTracker()
336 for c
in node
.children
[1:-1]:
337 if isinstance(c
, Leaf
):
340 for leaf
in c
.leaves():
343 return bt
.max_delimiter_priority()
349 def get_leaves_inside_matching_brackets(leaves
: Sequence
[Leaf
]) -> Set
[LeafID
]:
350 """Return leaves that are inside matching brackets.
352 The input `leaves` can have non-matching brackets at the head or tail parts.
353 Matching brackets are included.
356 # Start with the first opening bracket and ignore closing brackets before.
358 i
for i
, l
in enumerate(leaves
) if l
.type in OPENING_BRACKETS
360 except StopIteration:
364 for i
in range(start_index
, len(leaves
)):
366 if leaf
.type in OPENING_BRACKETS
:
367 bracket_stack
.append((BRACKET
[leaf
.type], i
))
368 if leaf
.type in CLOSING_BRACKETS
:
369 if bracket_stack
and leaf
.type == bracket_stack
[-1][0]:
370 _
, start
= bracket_stack
.pop()
371 for j
in range(start
, i
+ 1):
372 ids
.add(id(leaves
[j
]))