Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/black/nodes.py: 18%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2blib2to3 Node/Leaf transformation-related utility functions.
3"""
5from collections.abc import Iterator
6from typing import Final, Generic, Literal, TypeGuard, TypeVar, Union
8from mypy_extensions import mypyc_attr
10from black.mode import Mode, Preview
11from black.strings import get_string_prefix, has_triple_quotes
12from blib2to3 import pygram
13from blib2to3.pgen2 import token
14from blib2to3.pytree import NL, Leaf, Node, type_repr
16# The CLI cache directory is not known until Click parses the command line.
17pygram.initialize(save=False)
18syms: Final = pygram.python_symbols
21# types
22T = TypeVar("T")
23LN = Union[Leaf, Node]
24LeafID = int
25NodeType = int
28WHITESPACE: Final = {token.DEDENT, token.INDENT, token.NEWLINE}
29STATEMENT: Final = {
30 syms.if_stmt,
31 syms.while_stmt,
32 syms.for_stmt,
33 syms.try_stmt,
34 syms.except_clause,
35 syms.with_stmt,
36 syms.funcdef,
37 syms.classdef,
38 syms.match_stmt,
39 syms.case_block,
40}
41STANDALONE_COMMENT: Final = 153
42token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"
43LOGIC_OPERATORS: Final = {"and", "or"}
44COMPARATORS: Final = {
45 token.LESS,
46 token.GREATER,
47 token.EQEQUAL,
48 token.NOTEQUAL,
49 token.LESSEQUAL,
50 token.GREATEREQUAL,
51}
52MATH_OPERATORS: Final = {
53 token.VBAR,
54 token.CIRCUMFLEX,
55 token.AMPER,
56 token.LEFTSHIFT,
57 token.RIGHTSHIFT,
58 token.PLUS,
59 token.MINUS,
60 token.STAR,
61 token.SLASH,
62 token.DOUBLESLASH,
63 token.PERCENT,
64 token.AT,
65 token.TILDE,
66 token.DOUBLESTAR,
67}
68STARS: Final = {token.STAR, token.DOUBLESTAR}
69VARARGS_SPECIALS: Final = STARS | {token.SLASH}
70VARARGS_PARENTS: Final = {
71 syms.arglist,
72 syms.argument, # double star in arglist
73 syms.trailer, # single argument to call
74 syms.typedargslist,
75 syms.varargslist, # lambdas
76 syms.typevartuple, # star in a PEP 695 type parameter list
77 syms.paramspec, # double star in a PEP 695 type parameter list
78}
79UNPACKING_PARENTS: Final = {
80 syms.atom, # single element of a list or set literal
81 syms.dictsetmaker,
82 syms.listmaker,
83 syms.testlist_gexp,
84 syms.testlist_star_expr,
85 syms.subject_expr,
86 syms.pattern,
87}
88TEST_DESCENDANTS: Final = {
89 syms.test,
90 syms.lambdef,
91 syms.or_test,
92 syms.and_test,
93 syms.not_test,
94 syms.comparison,
95 syms.star_expr,
96 syms.expr,
97 syms.xor_expr,
98 syms.and_expr,
99 syms.shift_expr,
100 syms.arith_expr,
101 syms.trailer,
102 syms.term,
103 syms.power,
104 syms.namedexpr_test,
105}
106TYPED_NAMES: Final = {syms.tname, syms.tname_star}
107ASSIGNMENTS: Final = {
108 "=",
109 "+=",
110 "-=",
111 "*=",
112 "@=",
113 "/=",
114 "%=",
115 "&=",
116 "|=",
117 "^=",
118 "<<=",
119 ">>=",
120 "**=",
121 "//=",
122 ":",
123}
125IMPLICIT_TUPLE: Final = {syms.testlist, syms.testlist_star_expr, syms.exprlist}
126BRACKET: Final = {
127 token.LPAR: token.RPAR,
128 token.LSQB: token.RSQB,
129 token.LBRACE: token.RBRACE,
130}
131OPENING_BRACKETS: Final = set(BRACKET.keys())
132CLOSING_BRACKETS: Final = set(BRACKET.values())
133BRACKETS: Final = OPENING_BRACKETS | CLOSING_BRACKETS
134ALWAYS_NO_SPACE: Final = CLOSING_BRACKETS | {
135 token.COMMA,
136 STANDALONE_COMMENT,
137 token.FSTRING_MIDDLE,
138 token.FSTRING_END,
139 token.TSTRING_MIDDLE,
140 token.TSTRING_END,
141 token.BANG,
142}
145@mypyc_attr(allow_interpreted_subclasses=True)
146class Visitor(Generic[T]):
147 """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""
149 def visit(self, node: LN) -> Iterator[T]:
150 """Main method to visit `node` and its children.
152 It tries to find a `visit_*()` method for the given `node.type`, like
153 `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.
154 If no dedicated `visit_*()` method is found, chooses `visit_default()`
155 instead.
157 Then yields objects of type `T` from the selected visitor.
158 """
159 if node.type < 256:
160 name = token.tok_name[node.type]
161 else:
162 name = str(type_repr(node.type))
163 # We explicitly branch on whether a visitor exists (instead of
164 # using self.visit_default as the default arg to getattr) in order
165 # to save needing to create a bound method object and so mypyc can
166 # generate a native call to visit_default.
167 visitf = getattr(self, f"visit_{name}", None)
168 if visitf:
169 yield from visitf(node)
170 else:
171 yield from self.visit_default(node)
173 def visit_default(self, node: LN) -> Iterator[T]:
174 """Default `visit_*()` implementation. Recurses to children of `node`."""
175 if isinstance(node, Node):
176 for child in node.children:
177 yield from self.visit(child)
180def whitespace(leaf: Leaf, *, complex_subscript: bool, mode: Mode) -> str:
181 """Return whitespace prefix if needed for the given `leaf`.
183 `complex_subscript` signals whether the given leaf is part of a subscription
184 which has non-trivial arguments, like arithmetic expressions or function calls.
185 """
186 NO: Final[str] = ""
187 SPACE: Final[str] = " "
188 DOUBLESPACE: Final[str] = " "
189 t = leaf.type
190 p = leaf.parent
191 v = leaf.value
192 if t in ALWAYS_NO_SPACE:
193 return NO
195 if t == token.COMMENT:
196 return DOUBLESPACE
198 assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
199 if t == token.COLON and p.type not in {
200 syms.subscript,
201 syms.subscriptlist,
202 syms.sliceop,
203 }:
204 return NO
206 if t == token.LBRACE and p.type in (
207 syms.fstring_replacement_field,
208 syms.tstring_replacement_field,
209 ):
210 return NO
212 prev = leaf.prev_sibling
213 if not prev:
214 prevp = preceding_leaf(p)
215 if not prevp or prevp.type in OPENING_BRACKETS:
216 return NO
218 if t == token.COLON:
219 if prevp.type == token.COLON:
220 return NO
222 elif prevp.type != token.COMMA and not complex_subscript:
223 return NO
225 return SPACE
227 if prevp.type == token.EQUAL:
228 if prevp.parent:
229 if prevp.parent.type in {
230 syms.arglist,
231 syms.argument,
232 syms.parameters,
233 syms.varargslist,
234 }:
235 return NO
237 elif prevp.parent.type == syms.typedargslist:
238 # A bit hacky: if the equal sign has whitespace, it means we
239 # previously found it's a typed argument. So, we're using
240 # that, too.
241 return prevp.prefix
243 elif (
244 prevp.type == token.STAR
245 and parent_type(prevp) == syms.star_expr
246 and parent_type(prevp.parent) in (syms.subscriptlist, syms.tname_star)
247 ):
248 # No space between typevar tuples or unpacking them.
249 return NO
251 elif prevp.type in VARARGS_SPECIALS:
252 if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
253 return NO
255 elif prevp.type == token.COLON:
256 if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:
257 return SPACE if complex_subscript else NO
259 elif (
260 prevp.parent
261 and prevp.parent.type == syms.factor
262 and prevp.type in MATH_OPERATORS
263 ):
264 return NO
266 elif prevp.type == token.AT and p.parent and p.parent.type == syms.decorator:
267 # no space in decorators
268 return NO
270 elif prev.type in OPENING_BRACKETS:
271 return NO
273 elif prev.type == token.BANG:
274 return NO
276 if p.type in {syms.parameters, syms.arglist}:
277 # untyped function signatures or calls
278 if not prev or prev.type != token.COMMA:
279 return NO
281 elif p.type == syms.varargslist:
282 # lambdas
283 if prev and prev.type != token.COMMA:
284 return NO
286 elif p.type == syms.typedargslist:
287 # typed function signatures
288 if not prev:
289 return NO
291 if t == token.EQUAL:
292 if prev.type not in TYPED_NAMES:
293 return NO
295 elif prev.type == token.EQUAL:
296 # A bit hacky: if the equal sign has whitespace, it means we
297 # previously found it's a typed argument. So, we're using that, too.
298 return prev.prefix
300 elif prev.type != token.COMMA:
301 return NO
303 elif p.type in TYPED_NAMES:
304 # type names
305 if not prev:
306 prevp = preceding_leaf(p)
307 if not prevp or prevp.type != token.COMMA:
308 return NO
310 elif p.type == syms.trailer:
311 # attributes and calls
312 if t == token.LPAR or t == token.RPAR:
313 return NO
315 if not prev:
316 if t == token.DOT or t == token.LSQB:
317 return NO
319 elif prev.type != token.COMMA:
320 return NO
322 elif p.type == syms.argument:
323 # single argument
324 if t == token.EQUAL:
325 return NO
327 if not prev:
328 prevp = preceding_leaf(p)
329 if not prevp or prevp.type == token.LPAR:
330 return NO
332 elif prev.type in {token.EQUAL} | VARARGS_SPECIALS:
333 return NO
335 elif p.type == syms.decorator:
336 # decorators
337 return NO
339 elif p.type == syms.dotted_name:
340 if prev:
341 return NO
343 prevp = preceding_leaf(p)
344 if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
345 return NO
347 elif p.type == syms.classdef:
348 if t == token.LPAR:
349 return NO
351 if prev and prev.type == token.LPAR:
352 return NO
354 elif p.type in {syms.subscript, syms.sliceop}:
355 # indexing
356 if not prev:
357 assert p.parent is not None, "subscripts are always parented"
358 if p.parent.type == syms.subscriptlist:
359 return SPACE
361 return NO
363 elif t == token.COLONEQUAL or prev.type == token.COLONEQUAL:
364 return SPACE
366 elif not complex_subscript:
367 return NO
369 elif p.type == syms.atom:
370 if prev and t == token.DOT:
371 # dots, but not the first one.
372 return NO
374 elif p.type == syms.dictsetmaker:
375 # dict unpacking
376 if prev and prev.type == token.DOUBLESTAR:
377 return NO
379 elif p.type in {syms.factor, syms.star_expr}:
380 # unary ops
381 if not prev:
382 prevp = preceding_leaf(p)
383 if not prevp or prevp.type in OPENING_BRACKETS:
384 return NO
386 prevp_parent = prevp.parent
387 assert prevp_parent is not None
388 if prevp.type == token.COLON and prevp_parent.type in {
389 syms.subscript,
390 syms.sliceop,
391 }:
392 return NO
394 elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
395 return NO
397 elif t in {token.NAME, token.NUMBER, token.STRING}:
398 return NO
400 elif p.type == syms.import_from:
401 if t == token.DOT:
402 if prev and prev.type == token.DOT:
403 return NO
405 elif t == token.NAME:
406 if v == "import":
407 return SPACE
409 if prev and prev.type == token.DOT:
410 return NO
412 elif p.type == syms.sliceop:
413 return NO
415 elif p.type == syms.except_clause:
416 if t == token.STAR:
417 return NO
419 elif p.type == syms.typeparams:
420 # PEP 695 type parameter lists hug the preceding name
421 if t == token.LSQB:
422 return NO
424 elif p.type in {syms.typevartuple, syms.paramspec}:
425 # PEP 695 `*Ts` and `**P`
426 if prev and prev.type in STARS:
427 return NO
429 if Preview.simplify_power_operator_hugging in mode:
430 # Power operator hugging
431 if t == token.DOUBLESTAR and is_simple_exponentiation(p):
432 return NO
433 prevp = preceding_leaf(leaf)
434 if prevp and prevp.type == token.DOUBLESTAR:
435 if prevp.parent and is_simple_exponentiation(prevp.parent):
436 return NO
438 return SPACE
441def make_simple_prefix(nl_count: int, form_feed: bool, empty_line: str = "\n") -> str:
442 """Generate a normalized prefix string."""
443 if form_feed:
444 return (empty_line * (nl_count - 1)) + "\f" + empty_line
445 return empty_line * nl_count
448def preceding_leaf(node: LN | None) -> Leaf | None:
449 """Return the first leaf that precedes `node`, if any."""
450 while node:
451 res = node.prev_sibling
452 if res:
453 if isinstance(res, Leaf):
454 return res
456 try:
457 return list(res.leaves())[-1]
459 except IndexError:
460 return None
462 node = node.parent
463 return None
466def prev_siblings_are(node: LN | None, tokens: list[NodeType | None]) -> bool:
467 """Return if the `node` and its previous siblings match types against the provided
468 list of tokens; the provided `node` has its type matched against the last element in
469 the list. `None` can be used as the first element to declare that the start of the
470 list is anchored at the start of its parent's children."""
471 if not tokens:
472 return True
473 if tokens[-1] is None:
474 return node is None
475 if not node:
476 return False
477 if node.type != tokens[-1]:
478 return False
479 return prev_siblings_are(node.prev_sibling, tokens[:-1])
482def parent_type(node: LN | None) -> NodeType | None:
483 """
484 Returns:
485 @node.parent.type, if @node is not None and has a parent.
486 OR
487 None, otherwise.
488 """
489 if node is None or node.parent is None:
490 return None
492 return node.parent.type
495def child_towards(ancestor: Node, descendant: LN) -> LN | None:
496 """Return the child of `ancestor` that contains `descendant`."""
497 node: LN | None = descendant
498 while node and node.parent != ancestor:
499 node = node.parent
500 return node
503def replace_child(old_child: LN, new_child: LN) -> None:
504 """
505 Side Effects:
506 * If @old_child.parent is set, replace @old_child with @new_child in
507 @old_child's underlying Node structure.
508 OR
509 * Otherwise, this function does nothing.
510 """
511 parent = old_child.parent
512 if not parent:
513 return
515 child_idx = old_child.remove()
516 if child_idx is not None:
517 parent.insert_child(child_idx, new_child)
520def container_of(leaf: Leaf) -> LN:
521 """Return `leaf` or one of its ancestors that is the topmost container of it.
523 By "container" we mean a node where `leaf` is the very first child.
524 """
525 same_prefix = leaf.prefix
526 container: LN = leaf
527 while container:
528 parent = container.parent
529 if parent is None:
530 break
532 if parent.children[0].prefix != same_prefix:
533 break
535 if parent.type == syms.file_input:
536 break
538 if parent.prev_sibling is not None and parent.prev_sibling.type in BRACKETS:
539 break
541 container = parent
542 return container
545def first_leaf_of(node: LN) -> Leaf | None:
546 """Returns the first leaf of the node tree."""
547 if isinstance(node, Leaf):
548 return node
549 if node.children:
550 return first_leaf_of(node.children[0])
551 else:
552 return None
555def is_arith_like(node: LN) -> bool:
556 """Whether node is an arithmetic or a binary arithmetic expression"""
557 return node.type in {
558 syms.arith_expr,
559 syms.shift_expr,
560 syms.xor_expr,
561 syms.and_expr,
562 }
565def is_simple_exponentiation(node: LN) -> bool:
566 """Whether whitespace around `**` should be removed."""
568 def is_simple(node: LN) -> bool:
569 if isinstance(node, Leaf):
570 return node.type in (token.NAME, token.NUMBER, token.DOT, token.DOUBLESTAR)
571 elif node.type == syms.factor: # unary operators
572 return is_simple(node.children[1])
573 else:
574 return all(is_simple(child) for child in node.children)
576 return (
577 node.type == syms.power
578 and len(node.children) >= 3
579 and node.children[-2].type == token.DOUBLESTAR
580 and is_simple(node)
581 )
584def is_docstring(node: NL) -> bool:
585 if isinstance(node, Leaf):
586 if node.type != token.STRING:
587 return False
589 prefix = get_string_prefix(node.value)
590 # Bytes, f-strings and t-strings never evaluate to str, so they are not
591 # docstrings even in docstring position.
592 if set(prefix).intersection("bBfFtT"):
593 return False
595 if (
596 node.parent
597 and node.parent.type == syms.simple_stmt
598 and not node.parent.prev_sibling
599 and node.parent.parent
600 and node.parent.parent.type == syms.file_input
601 ):
602 return True
604 if prev_siblings_are(
605 node.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
606 ):
607 return True
609 # Multiline docstring on the same line as the `def`.
610 if prev_siblings_are(node.parent, [syms.parameters, token.COLON, syms.simple_stmt]):
611 # `syms.parameters` is only used in funcdefs and async_funcdefs in the Python
612 # grammar. We're safe to return True without further checks.
613 return True
615 return False
618def is_empty_tuple(node: LN) -> bool:
619 """Return True if `node` holds an empty tuple."""
620 return (
621 node.type == syms.atom
622 and len(node.children) == 2
623 and node.children[0].type == token.LPAR
624 and node.children[1].type == token.RPAR
625 )
628def is_one_tuple(node: LN) -> bool:
629 """Return True if `node` holds a tuple with one element, with or without parens."""
630 if node.type == syms.atom:
631 gexp = unwrap_singleton_parenthesis(node)
632 if gexp is None or gexp.type != syms.testlist_gexp:
633 return False
635 return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA
637 return (
638 node.type in IMPLICIT_TUPLE
639 and len(node.children) == 2
640 and node.children[1].type == token.COMMA
641 )
644def is_one_element_subscript(node: LN) -> bool:
645 """Return True if `node` is a subscript trailer with one element, like `[x,]`."""
646 return (
647 node.type == syms.trailer
648 and node.children[1].type == syms.subscriptlist
649 and len(node.children[1].children) == 2
650 and node.children[1].children[1].type == token.COMMA
651 )
654def is_tuple(node: LN) -> bool:
655 """Return True if `node` holds a tuple."""
656 if node.type != syms.atom:
657 return False
658 gexp = unwrap_singleton_parenthesis(node)
659 if gexp is None or gexp.type != syms.testlist_gexp:
660 return False
662 return True
665def is_list(node: LN) -> bool:
666 """Return True if `node` holds a list or list comprehension."""
667 return (
668 node.type == syms.atom
669 and len(node.children) >= 2
670 and node.children[0].type == token.LSQB
671 and node.children[-1].type == token.RSQB
672 )
675def is_tuple_containing_walrus(node: LN) -> bool:
676 """Return True if `node` holds a tuple that contains a walrus operator."""
677 if node.type != syms.atom:
678 return False
679 gexp = unwrap_singleton_parenthesis(node)
680 if gexp is None or gexp.type != syms.testlist_gexp:
681 return False
683 return any(child.type == syms.namedexpr_test for child in gexp.children)
686def is_tuple_containing_star(node: LN) -> bool:
687 """Return True if `node` holds a tuple that contains a star operator."""
688 if node.type != syms.atom:
689 return False
690 gexp = unwrap_singleton_parenthesis(node)
691 if gexp is None or gexp.type != syms.testlist_gexp:
692 return False
694 return any(child.type == syms.star_expr for child in gexp.children)
697def is_generator(node: LN) -> bool:
698 """Return True if `node` holds a generator."""
699 if node.type != syms.atom:
700 return False
701 gexp = unwrap_singleton_parenthesis(node)
702 if gexp is None or gexp.type != syms.testlist_gexp:
703 return False
705 return any(child.type == syms.old_comp_for for child in gexp.children)
708def is_one_sequence_between(
709 opening: Leaf,
710 closing: Leaf,
711 leaves: list[Leaf],
712 brackets: tuple[int, int] = (token.LPAR, token.RPAR),
713) -> bool:
714 """Return True if content between `opening` and `closing` is a one-sequence."""
715 if (opening.type, closing.type) != brackets:
716 return False
718 depth = closing.bracket_depth + 1
719 # Locate `opening` by scanning inward from both ends at once. Callers pass the
720 # whole line's leaf list and this runs once per bracket, so a plain forward scan
721 # from the start costs O(index) and turns quadratic on a long line; meeting in
722 # the middle bounds each lookup to the nearer end.
723 _opening_index = -1
724 left = 0
725 right = len(leaves) - 1
726 while left <= right:
727 if leaves[left] is opening:
728 _opening_index = left
729 break
730 if leaves[right] is opening:
731 _opening_index = right
732 break
733 left += 1
734 right -= 1
736 if _opening_index == -1:
737 return False
739 commas = 0
740 _opening_index += 1
741 for leaf in leaves[_opening_index:]:
742 if leaf is closing:
743 break
745 bracket_depth = leaf.bracket_depth
746 if bracket_depth == depth and leaf.type == token.COMMA:
747 commas += 1
748 if leaf.parent and leaf.parent.type in {
749 syms.arglist,
750 syms.typedargslist,
751 }:
752 commas += 1
753 break
755 return commas < 2
758def is_walrus_assignment(node: LN) -> bool:
759 """Return True iff `node` is of the shape ( test := test )"""
760 inner = unwrap_singleton_parenthesis(node)
761 return inner is not None and inner.type == syms.namedexpr_test
764def is_simple_decorator_trailer(node: LN, last: bool = False) -> bool:
765 """Return True iff `node` is a trailer valid in a simple decorator"""
766 return node.type == syms.trailer and (
767 (
768 len(node.children) == 2
769 and node.children[0].type == token.DOT
770 and node.children[1].type == token.NAME
771 )
772 # last trailer can be an argument-less parentheses pair
773 or (
774 last
775 and len(node.children) == 2
776 and node.children[0].type == token.LPAR
777 and node.children[1].type == token.RPAR
778 )
779 # last trailer can be arguments
780 or (
781 last
782 and len(node.children) == 3
783 and node.children[0].type == token.LPAR
784 # and node.children[1].type == syms.argument
785 and node.children[2].type == token.RPAR
786 )
787 )
790def is_simple_decorator_expression(node: LN) -> bool:
791 """Return True iff `node` could be a 'dotted name' decorator
793 This function takes the node of the 'namedexpr_test' of the new decorator
794 grammar and test if it would be valid under the old decorator grammar.
796 The old grammar was: decorator: @ dotted_name [arguments] NEWLINE
797 The new grammar is : decorator: @ namedexpr_test NEWLINE
798 """
799 if node.type == token.NAME:
800 return True
801 if node.type == syms.power:
802 if node.children:
803 return (
804 node.children[0].type == token.NAME
805 and all(map(is_simple_decorator_trailer, node.children[1:-1]))
806 and (
807 len(node.children) < 2
808 or is_simple_decorator_trailer(node.children[-1], last=True)
809 )
810 )
811 return False
814def is_yield(node: LN) -> bool:
815 """Return True if `node` holds a `yield` or `yield from` expression."""
816 if node.type == syms.yield_expr:
817 return True
819 if is_name_token(node) and node.value == "yield":
820 return True
822 if node.type != syms.atom:
823 return False
825 if len(node.children) != 3:
826 return False
828 lpar, expr, rpar = node.children
829 if lpar.type == token.LPAR and rpar.type == token.RPAR:
830 return is_yield(expr)
832 return False
835def is_vararg(leaf: Leaf, within: set[NodeType]) -> bool:
836 """Return True if `leaf` is a star or double star in a vararg or kwarg.
838 If `within` includes VARARGS_PARENTS, this applies to function signatures.
839 If `within` includes UNPACKING_PARENTS, it applies to right hand-side
840 extended iterable unpacking (PEP 3132) and additional unpacking
841 generalizations (PEP 448).
842 """
843 if leaf.type not in VARARGS_SPECIALS or not leaf.parent:
844 return False
846 p = leaf.parent
847 if p.type == syms.star_expr:
848 # Star expressions are also used as assignment targets in extended
849 # iterable unpacking (PEP 3132). See what its parent is instead.
850 if not p.parent:
851 return False
853 p = p.parent
855 return p.type in within
858def is_fstring(node: Node) -> bool:
859 """Return True if the node is an f-string"""
860 return node.type == syms.fstring
863def fstring_tstring_to_string(node: Node) -> Leaf:
864 """Converts an fstring or tstring node back to a string node."""
865 string_without_prefix = str(node)[len(node.prefix) :]
866 string_leaf = Leaf(token.STRING, string_without_prefix, prefix=node.prefix)
867 string_leaf.lineno = node.get_lineno() or 0
868 return string_leaf
871def is_multiline_string(node: LN) -> bool:
872 """Return True if `leaf` is a multiline string that actually spans many lines."""
873 if isinstance(node, Node) and is_fstring(node):
874 leaf = fstring_tstring_to_string(node)
875 elif isinstance(node, Leaf):
876 leaf = node
877 else:
878 return False
880 return has_triple_quotes(leaf.value) and "\n" in leaf.value
883def is_parent_function_or_class(node: Node) -> bool:
884 assert node.type in {syms.suite, syms.simple_stmt}
885 assert node.parent is not None
886 # Note this works for suites / simple_stmts in async def as well
887 return node.parent.type in {syms.funcdef, syms.classdef}
890def is_stub_suite(node: Node) -> bool:
891 """Return True if `node` is a suite with a stub body."""
892 if node.parent is not None and not is_parent_function_or_class(node):
893 return False
895 # If there is a comment, we want to keep it.
896 if node.prefix.strip():
897 return False
899 if (
900 len(node.children) != 4
901 or node.children[0].type != token.NEWLINE
902 or node.children[1].type != token.INDENT
903 or node.children[3].type != token.DEDENT
904 ):
905 return False
907 if node.children[3].prefix.strip():
908 return False
910 return is_stub_body(node.children[2])
913def is_stub_body(node: LN) -> bool:
914 """Return True if `node` is a simple statement containing an ellipsis."""
915 if not isinstance(node, Node) or node.type != syms.simple_stmt:
916 return False
918 if len(node.children) != 2:
919 return False
921 child = node.children[0]
922 return (
923 not child.prefix.strip()
924 and child.type == syms.atom
925 and len(child.children) == 3
926 and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
927 )
930def is_atom_with_invisible_parens(node: LN) -> bool:
931 """Given a `LN`, determines whether it's an atom `node` with invisible
932 parens. Useful in dedupe-ing and normalizing parens.
933 """
934 if isinstance(node, Leaf) or node.type != syms.atom:
935 return False
937 first, last = node.children[0], node.children[-1]
938 return (
939 isinstance(first, Leaf)
940 and first.type == token.LPAR
941 and first.value == ""
942 and isinstance(last, Leaf)
943 and last.type == token.RPAR
944 and last.value == ""
945 )
948def is_empty_par(leaf: Leaf) -> bool:
949 return is_empty_lpar(leaf) or is_empty_rpar(leaf)
952def is_empty_lpar(leaf: Leaf) -> bool:
953 return leaf.type == token.LPAR and leaf.value == ""
956def is_empty_rpar(leaf: Leaf) -> bool:
957 return leaf.type == token.RPAR and leaf.value == ""
960def is_import(leaf: Leaf) -> bool:
961 """Return True if the given leaf starts an import statement."""
962 p = leaf.parent
963 t = leaf.type
964 v = leaf.value
965 return bool(
966 (t == token.LAZY and p and p.type == syms.lazy_import)
967 or (
968 t == token.NAME
969 and (
970 (v == "import" and p and p.type == syms.import_name)
971 or (v == "from" and p and p.type == syms.import_from)
972 )
973 )
974 )
977def is_with_or_async_with_stmt(leaf: Leaf) -> bool:
978 """Return True if the given leaf starts a with or async with statement."""
979 return bool(
980 leaf.type == token.NAME
981 and leaf.value == "with"
982 and leaf.parent
983 and leaf.parent.type == syms.with_stmt
984 ) or bool(
985 leaf.type == token.ASYNC
986 and leaf.next_sibling
987 and leaf.next_sibling.type == syms.with_stmt
988 )
991def is_async_stmt_or_funcdef(leaf: Leaf) -> bool:
992 """Return True if the given leaf starts an async def/for/with statement.
994 Note that `async def` can be either an `async_stmt` or `async_funcdef`,
995 the latter is used when it has decorators.
996 """
997 return bool(
998 leaf.type == token.ASYNC
999 and leaf.parent
1000 and leaf.parent.type in {syms.async_stmt, syms.async_funcdef}
1001 )
1004def is_type_comment(leaf: Leaf, mode: Mode) -> bool:
1005 """Return True if the given leaf is a type comment. This function should only
1006 be used for general type comments (excluding ignore annotations, which should
1007 use `is_type_ignore_comment`). Note that general type comments are no longer
1008 used in modern version of Python, this function may be deprecated in the future."""
1009 t = leaf.type
1010 v = leaf.value
1011 return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_comment_string(v, mode)
1014def is_type_comment_string(value: str, mode: Mode) -> bool:
1015 return value.startswith("#") and value[1:].lstrip().startswith("type:")
1018def is_type_ignore_comment(leaf: Leaf, mode: Mode) -> bool:
1019 """Return True if the given leaf is a type comment with ignore annotation."""
1020 t = leaf.type
1021 v = leaf.value
1022 return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_ignore_comment_string(
1023 v, mode
1024 )
1027def is_type_ignore_comment_string(value: str, mode: Mode) -> bool:
1028 """Return True if the given string match with type comment with
1029 ignore annotation."""
1030 return is_type_comment_string(value, mode) and value.split(":", 1)[
1031 1
1032 ].lstrip().startswith("ignore")
1035def wrap_in_parentheses(
1036 parent: Node, child: LN, *, visible: bool = True, index: int | None = None
1037) -> None:
1038 """Wrap `child` in parentheses.
1040 This replaces `child` with an atom holding the parentheses and the old
1041 child. That requires moving the prefix.
1043 If `visible` is False, the leaves will be valueless (and thus invisible).
1045 When the caller already knows `child`'s position in `parent.children` it
1046 can pass `index` so the child is swapped in place. The default path locates
1047 `child` with `Base.remove`, which scans and rewrites the whole child list;
1048 that is O(len(parent.children)) per call and turns quadratic when a caller
1049 wraps every child of a large node (e.g. each value of a big dict literal).
1050 """
1051 lpar = Leaf(token.LPAR, "(" if visible else "")
1052 rpar = Leaf(token.RPAR, ")" if visible else "")
1053 prefix = child.prefix
1054 child.prefix = ""
1055 if index is None:
1056 index = child.remove() or 0
1057 new_child = Node(syms.atom, [lpar, child, rpar])
1058 new_child.prefix = prefix
1059 parent.insert_child(index, new_child)
1060 else:
1061 # Detach the child pointer first so the atom can adopt it, then swap it
1062 # in place. set_child resets the old child's parent, so reattach it to
1063 # the new atom afterwards.
1064 child.parent = None
1065 new_child = Node(syms.atom, [lpar, child, rpar])
1066 new_child.prefix = prefix
1067 parent.set_child(index, new_child)
1068 child.parent = new_child
1071def unwrap_singleton_parenthesis(node: LN) -> LN | None:
1072 """Returns `wrapped` if `node` is of the shape ( wrapped ).
1074 Parenthesis can be optional. Returns None otherwise"""
1075 if len(node.children) != 3:
1076 return None
1078 lpar, wrapped, rpar = node.children
1079 if not (lpar.type == token.LPAR and rpar.type == token.RPAR):
1080 return None
1082 return wrapped
1085def ensure_visible(leaf: Leaf) -> None:
1086 """Make sure parentheses are visible.
1088 They could be invisible as part of some statements (see
1089 :func:`normalize_invisible_parens` and :func:`visit_import_from`).
1090 """
1091 if leaf.type == token.LPAR:
1092 leaf.value = "("
1093 elif leaf.type == token.RPAR:
1094 leaf.value = ")"
1097def is_name_token(nl: NL) -> TypeGuard[Leaf]:
1098 return nl.type == token.NAME
1101def is_lpar_token(nl: NL) -> TypeGuard[Leaf]:
1102 return nl.type == token.LPAR
1105def is_rpar_token(nl: NL) -> TypeGuard[Leaf]:
1106 return nl.type == token.RPAR
1109def is_number_token(nl: NL) -> TypeGuard[Leaf]:
1110 return nl.type == token.NUMBER
1113def get_annotation_type(leaf: Leaf) -> Literal["return", "param", None]:
1114 """Returns the type of annotation this leaf is part of, if any."""
1115 ancestor = leaf.parent
1116 while ancestor is not None:
1117 if ancestor.prev_sibling and ancestor.prev_sibling.type == token.RARROW:
1118 return "return"
1119 if ancestor.parent and ancestor.parent.type == syms.tname:
1120 return "param"
1121 ancestor = ancestor.parent
1122 return None
1125def is_part_of_annotation(leaf: Leaf) -> bool:
1126 """Returns whether this leaf is part of a type annotation."""
1127 assert leaf.parent is not None
1128 return get_annotation_type(leaf) is not None
1131def first_leaf(node: LN) -> Leaf | None:
1132 """Returns the first leaf of the ancestor node."""
1133 if isinstance(node, Leaf):
1134 return node
1135 elif not node.children:
1136 return None
1137 else:
1138 return first_leaf(node.children[0])
1141def last_leaf(node: LN) -> Leaf | None:
1142 """Returns the last leaf of the ancestor node."""
1143 if isinstance(node, Leaf):
1144 return node
1145 elif not node.children:
1146 return None
1147 else:
1148 return last_leaf(node.children[-1])
1151def furthest_ancestor_with_last_leaf(leaf: Leaf) -> LN:
1152 """Returns the furthest ancestor that has this leaf node as the last leaf."""
1153 node: LN = leaf
1154 while node.parent and node.parent.children and node is node.parent.children[-1]:
1155 node = node.parent
1156 return node
1159def has_sibling_with_type(node: LN, type: int) -> bool:
1160 # Check previous siblings
1161 sibling = node.prev_sibling
1162 while sibling is not None:
1163 if sibling.type == type:
1164 return True
1165 sibling = sibling.prev_sibling
1167 # Check next siblings
1168 sibling = node.next_sibling
1169 while sibling is not None:
1170 if sibling.type == type:
1171 return True
1172 sibling = sibling.next_sibling
1174 return False