]> git.madduck.net Git - etc/vim.git/blob - src/black/brackets.py

madduck's git repository

Every one of the projects in this repository is available at the canonical URL git://git.madduck.net/madduck/pub/<projectpath> — see each project's metadata for the exact URL.

All patches and comments are welcome. Please squash your changes to logical commits before using git-format-patch and git-send-email to patches@git.madduck.net. If you'd read over the Git project's submission guidelines and adhered to them, I'd be especially grateful.

SSH access, as well as push access can be individually arranged.

If you use my repositories frequently, consider adding the following snippet to ~/.gitconfig and using the third clone URL listed for each project:

[url "git://git.madduck.net/madduck/"]
  insteadOf = madduck:

Make README logo link to docs (#3285)
[etc/vim.git] / 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, 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