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.cache import CACHE_DIR
11from black.mode import Mode, Preview
12from black.strings import get_string_prefix, has_triple_quotes
13from blib2to3 import pygram
14from blib2to3.pgen2 import token
15from blib2to3.pytree import NL, Leaf, Node, type_repr
17pygram.initialize(CACHE_DIR)
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}
144RARROW = 55
147@mypyc_attr(allow_interpreted_subclasses=True)
148class Visitor(Generic[T]):
149 """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""
151 def visit(self, node: LN) -> Iterator[T]:
152 """Main method to visit `node` and its children.
154 It tries to find a `visit_*()` method for the given `node.type`, like
155 `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.
156 If no dedicated `visit_*()` method is found, chooses `visit_default()`
157 instead.
159 Then yields objects of type `T` from the selected visitor.
160 """
161 if node.type < 256:
162 name = token.tok_name[node.type]
163 else:
164 name = str(type_repr(node.type))
165 # We explicitly branch on whether a visitor exists (instead of
166 # using self.visit_default as the default arg to getattr) in order
167 # to save needing to create a bound method object and so mypyc can
168 # generate a native call to visit_default.
169 visitf = getattr(self, f"visit_{name}", None)
170 if visitf:
171 yield from visitf(node)
172 else:
173 yield from self.visit_default(node)
175 def visit_default(self, node: LN) -> Iterator[T]:
176 """Default `visit_*()` implementation. Recurses to children of `node`."""
177 if isinstance(node, Node):
178 for child in node.children:
179 yield from self.visit(child)
182def whitespace(leaf: Leaf, *, complex_subscript: bool, mode: Mode) -> str:
183 """Return whitespace prefix if needed for the given `leaf`.
185 `complex_subscript` signals whether the given leaf is part of a subscription
186 which has non-trivial arguments, like arithmetic expressions or function calls.
187 """
188 NO: Final[str] = ""
189 SPACE: Final[str] = " "
190 DOUBLESPACE: Final[str] = " "
191 t = leaf.type
192 p = leaf.parent
193 v = leaf.value
194 if t in ALWAYS_NO_SPACE:
195 return NO
197 if t == token.COMMENT:
198 return DOUBLESPACE
200 assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
201 if t == token.COLON and p.type not in {
202 syms.subscript,
203 syms.subscriptlist,
204 syms.sliceop,
205 }:
206 return NO
208 if t == token.LBRACE and p.type in (
209 syms.fstring_replacement_field,
210 syms.tstring_replacement_field,
211 ):
212 return NO
214 prev = leaf.prev_sibling
215 if not prev:
216 prevp = preceding_leaf(p)
217 if not prevp or prevp.type in OPENING_BRACKETS:
218 return NO
220 if t == token.COLON:
221 if prevp.type == token.COLON:
222 return NO
224 elif prevp.type != token.COMMA and not complex_subscript:
225 return NO
227 return SPACE
229 if prevp.type == token.EQUAL:
230 if prevp.parent:
231 if prevp.parent.type in {
232 syms.arglist,
233 syms.argument,
234 syms.parameters,
235 syms.varargslist,
236 }:
237 return NO
239 elif prevp.parent.type == syms.typedargslist:
240 # A bit hacky: if the equal sign has whitespace, it means we
241 # previously found it's a typed argument. So, we're using
242 # that, too.
243 return prevp.prefix
245 elif (
246 prevp.type == token.STAR
247 and parent_type(prevp) == syms.star_expr
248 and parent_type(prevp.parent) in (syms.subscriptlist, syms.tname_star)
249 ):
250 # No space between typevar tuples or unpacking them.
251 return NO
253 elif prevp.type in VARARGS_SPECIALS:
254 if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
255 return NO
257 elif prevp.type == token.COLON:
258 if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:
259 return SPACE if complex_subscript else NO
261 elif (
262 prevp.parent
263 and prevp.parent.type == syms.factor
264 and prevp.type in MATH_OPERATORS
265 ):
266 return NO
268 elif prevp.type == token.AT and p.parent and p.parent.type == syms.decorator:
269 # no space in decorators
270 return NO
272 elif prev.type in OPENING_BRACKETS:
273 return NO
275 elif prev.type == token.BANG:
276 return NO
278 if p.type in {syms.parameters, syms.arglist}:
279 # untyped function signatures or calls
280 if not prev or prev.type != token.COMMA:
281 return NO
283 elif p.type == syms.varargslist:
284 # lambdas
285 if prev and prev.type != token.COMMA:
286 return NO
288 elif p.type == syms.typedargslist:
289 # typed function signatures
290 if not prev:
291 return NO
293 if t == token.EQUAL:
294 if prev.type not in TYPED_NAMES:
295 return NO
297 elif prev.type == token.EQUAL:
298 # A bit hacky: if the equal sign has whitespace, it means we
299 # previously found it's a typed argument. So, we're using that, too.
300 return prev.prefix
302 elif prev.type != token.COMMA:
303 return NO
305 elif p.type in TYPED_NAMES:
306 # type names
307 if not prev:
308 prevp = preceding_leaf(p)
309 if not prevp or prevp.type != token.COMMA:
310 return NO
312 elif p.type == syms.trailer:
313 # attributes and calls
314 if t == token.LPAR or t == token.RPAR:
315 return NO
317 if not prev:
318 if t == token.DOT or t == token.LSQB:
319 return NO
321 elif prev.type != token.COMMA:
322 return NO
324 elif p.type == syms.argument:
325 # single argument
326 if t == token.EQUAL:
327 return NO
329 if not prev:
330 prevp = preceding_leaf(p)
331 if not prevp or prevp.type == token.LPAR:
332 return NO
334 elif prev.type in {token.EQUAL} | VARARGS_SPECIALS:
335 return NO
337 elif p.type == syms.decorator:
338 # decorators
339 return NO
341 elif p.type == syms.dotted_name:
342 if prev:
343 return NO
345 prevp = preceding_leaf(p)
346 if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
347 return NO
349 elif p.type == syms.classdef:
350 if t == token.LPAR:
351 return NO
353 if prev and prev.type == token.LPAR:
354 return NO
356 elif p.type in {syms.subscript, syms.sliceop}:
357 # indexing
358 if not prev:
359 assert p.parent is not None, "subscripts are always parented"
360 if p.parent.type == syms.subscriptlist:
361 return SPACE
363 return NO
365 elif t == token.COLONEQUAL or prev.type == token.COLONEQUAL:
366 return SPACE
368 elif not complex_subscript:
369 return NO
371 elif p.type == syms.atom:
372 if prev and t == token.DOT:
373 # dots, but not the first one.
374 return NO
376 elif p.type == syms.dictsetmaker:
377 # dict unpacking
378 if prev and prev.type == token.DOUBLESTAR:
379 return NO
381 elif p.type in {syms.factor, syms.star_expr}:
382 # unary ops
383 if not prev:
384 prevp = preceding_leaf(p)
385 if not prevp or prevp.type in OPENING_BRACKETS:
386 return NO
388 prevp_parent = prevp.parent
389 assert prevp_parent is not None
390 if prevp.type == token.COLON and prevp_parent.type in {
391 syms.subscript,
392 syms.sliceop,
393 }:
394 return NO
396 elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
397 return NO
399 elif t in {token.NAME, token.NUMBER, token.STRING}:
400 return NO
402 elif p.type == syms.import_from:
403 if t == token.DOT:
404 if prev and prev.type == token.DOT:
405 return NO
407 elif t == token.NAME:
408 if v == "import":
409 return SPACE
411 if prev and prev.type == token.DOT:
412 return NO
414 elif p.type == syms.sliceop:
415 return NO
417 elif p.type == syms.except_clause:
418 if t == token.STAR:
419 return NO
421 if Preview.simplify_power_operator_hugging in mode:
422 # Power operator hugging
423 if t == token.DOUBLESTAR and is_simple_exponentiation(p):
424 return NO
425 prevp = preceding_leaf(leaf)
426 if prevp and prevp.type == token.DOUBLESTAR:
427 if prevp.parent and is_simple_exponentiation(prevp.parent):
428 return NO
430 return SPACE
433def make_simple_prefix(nl_count: int, form_feed: bool, empty_line: str = "\n") -> str:
434 """Generate a normalized prefix string."""
435 if form_feed:
436 return (empty_line * (nl_count - 1)) + "\f" + empty_line
437 return empty_line * nl_count
440def preceding_leaf(node: LN | None) -> Leaf | None:
441 """Return the first leaf that precedes `node`, if any."""
442 while node:
443 res = node.prev_sibling
444 if res:
445 if isinstance(res, Leaf):
446 return res
448 try:
449 return list(res.leaves())[-1]
451 except IndexError:
452 return None
454 node = node.parent
455 return None
458def prev_siblings_are(node: LN | None, tokens: list[NodeType | None]) -> bool:
459 """Return if the `node` and its previous siblings match types against the provided
460 list of tokens; the provided `node`has its type matched against the last element in
461 the list. `None` can be used as the first element to declare that the start of the
462 list is anchored at the start of its parent's children."""
463 if not tokens:
464 return True
465 if tokens[-1] is None:
466 return node is None
467 if not node:
468 return False
469 if node.type != tokens[-1]:
470 return False
471 return prev_siblings_are(node.prev_sibling, tokens[:-1])
474def parent_type(node: LN | None) -> NodeType | None:
475 """
476 Returns:
477 @node.parent.type, if @node is not None and has a parent.
478 OR
479 None, otherwise.
480 """
481 if node is None or node.parent is None:
482 return None
484 return node.parent.type
487def child_towards(ancestor: Node, descendant: LN) -> LN | None:
488 """Return the child of `ancestor` that contains `descendant`."""
489 node: LN | None = descendant
490 while node and node.parent != ancestor:
491 node = node.parent
492 return node
495def replace_child(old_child: LN, new_child: LN) -> None:
496 """
497 Side Effects:
498 * If @old_child.parent is set, replace @old_child with @new_child in
499 @old_child's underlying Node structure.
500 OR
501 * Otherwise, this function does nothing.
502 """
503 parent = old_child.parent
504 if not parent:
505 return
507 child_idx = old_child.remove()
508 if child_idx is not None:
509 parent.insert_child(child_idx, new_child)
512def container_of(leaf: Leaf) -> LN:
513 """Return `leaf` or one of its ancestors that is the topmost container of it.
515 By "container" we mean a node where `leaf` is the very first child.
516 """
517 same_prefix = leaf.prefix
518 container: LN = leaf
519 while container:
520 parent = container.parent
521 if parent is None:
522 break
524 if parent.children[0].prefix != same_prefix:
525 break
527 if parent.type == syms.file_input:
528 break
530 if parent.prev_sibling is not None and parent.prev_sibling.type in BRACKETS:
531 break
533 container = parent
534 return container
537def first_leaf_of(node: LN) -> Leaf | None:
538 """Returns the first leaf of the node tree."""
539 if isinstance(node, Leaf):
540 return node
541 if node.children:
542 return first_leaf_of(node.children[0])
543 else:
544 return None
547def is_arith_like(node: LN) -> bool:
548 """Whether node is an arithmetic or a binary arithmetic expression"""
549 return node.type in {
550 syms.arith_expr,
551 syms.shift_expr,
552 syms.xor_expr,
553 syms.and_expr,
554 }
557def is_simple_exponentiation(node: LN) -> bool:
558 """Whether whitespace around `**` should be removed."""
560 def is_simple(node: LN) -> bool:
561 if isinstance(node, Leaf):
562 return node.type in (token.NAME, token.NUMBER, token.DOT, token.DOUBLESTAR)
563 elif node.type == syms.factor: # unary operators
564 return is_simple(node.children[1])
565 else:
566 return all(is_simple(child) for child in node.children)
568 return (
569 node.type == syms.power
570 and len(node.children) >= 3
571 and node.children[-2].type == token.DOUBLESTAR
572 and is_simple(node)
573 )
576def is_docstring(node: NL) -> bool:
577 if isinstance(node, Leaf):
578 if node.type != token.STRING:
579 return False
581 prefix = get_string_prefix(node.value)
582 if set(prefix).intersection("bBfF"):
583 return False
585 if (
586 node.parent
587 and node.parent.type == syms.simple_stmt
588 and not node.parent.prev_sibling
589 and node.parent.parent
590 and node.parent.parent.type == syms.file_input
591 ):
592 return True
594 if prev_siblings_are(
595 node.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
596 ):
597 return True
599 # Multiline docstring on the same line as the `def`.
600 if prev_siblings_are(node.parent, [syms.parameters, token.COLON, syms.simple_stmt]):
601 # `syms.parameters` is only used in funcdefs and async_funcdefs in the Python
602 # grammar. We're safe to return True without further checks.
603 return True
605 return False
608def is_empty_tuple(node: LN) -> bool:
609 """Return True if `node` holds an empty tuple."""
610 return (
611 node.type == syms.atom
612 and len(node.children) == 2
613 and node.children[0].type == token.LPAR
614 and node.children[1].type == token.RPAR
615 )
618def is_one_tuple(node: LN) -> bool:
619 """Return True if `node` holds a tuple with one element, with or without parens."""
620 if node.type == syms.atom:
621 gexp = unwrap_singleton_parenthesis(node)
622 if gexp is None or gexp.type != syms.testlist_gexp:
623 return False
625 return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA
627 return (
628 node.type in IMPLICIT_TUPLE
629 and len(node.children) == 2
630 and node.children[1].type == token.COMMA
631 )
634def is_tuple(node: LN) -> bool:
635 """Return True if `node` holds a tuple."""
636 if node.type != syms.atom:
637 return False
638 gexp = unwrap_singleton_parenthesis(node)
639 if gexp is None or gexp.type != syms.testlist_gexp:
640 return False
642 return True
645def is_tuple_containing_walrus(node: LN) -> bool:
646 """Return True if `node` holds a tuple that contains a walrus operator."""
647 if node.type != syms.atom:
648 return False
649 gexp = unwrap_singleton_parenthesis(node)
650 if gexp is None or gexp.type != syms.testlist_gexp:
651 return False
653 return any(child.type == syms.namedexpr_test for child in gexp.children)
656def is_tuple_containing_star(node: LN) -> bool:
657 """Return True if `node` holds a tuple that contains a star operator."""
658 if node.type != syms.atom:
659 return False
660 gexp = unwrap_singleton_parenthesis(node)
661 if gexp is None or gexp.type != syms.testlist_gexp:
662 return False
664 return any(child.type == syms.star_expr for child in gexp.children)
667def is_generator(node: LN) -> bool:
668 """Return True if `node` holds a generator."""
669 if node.type != syms.atom:
670 return False
671 gexp = unwrap_singleton_parenthesis(node)
672 if gexp is None or gexp.type != syms.testlist_gexp:
673 return False
675 return any(child.type == syms.old_comp_for for child in gexp.children)
678def is_one_sequence_between(
679 opening: Leaf,
680 closing: Leaf,
681 leaves: list[Leaf],
682 brackets: tuple[int, int] = (token.LPAR, token.RPAR),
683) -> bool:
684 """Return True if content between `opening` and `closing` is a one-sequence."""
685 if (opening.type, closing.type) != brackets:
686 return False
688 depth = closing.bracket_depth + 1
689 # Locate `opening` by scanning inward from both ends at once. Callers pass the
690 # whole line's leaf list and this runs once per bracket, so a plain forward scan
691 # from the start costs O(index) and turns quadratic on a long line; meeting in
692 # the middle bounds each lookup to the nearer end.
693 _opening_index = -1
694 left = 0
695 right = len(leaves) - 1
696 while left <= right:
697 if leaves[left] is opening:
698 _opening_index = left
699 break
700 if leaves[right] is opening:
701 _opening_index = right
702 break
703 left += 1
704 right -= 1
706 if _opening_index == -1:
707 return False
709 commas = 0
710 _opening_index += 1
711 for leaf in leaves[_opening_index:]:
712 if leaf is closing:
713 break
715 bracket_depth = leaf.bracket_depth
716 if bracket_depth == depth and leaf.type == token.COMMA:
717 commas += 1
718 if leaf.parent and leaf.parent.type in {
719 syms.arglist,
720 syms.typedargslist,
721 }:
722 commas += 1
723 break
725 return commas < 2
728def is_walrus_assignment(node: LN) -> bool:
729 """Return True iff `node` is of the shape ( test := test )"""
730 inner = unwrap_singleton_parenthesis(node)
731 return inner is not None and inner.type == syms.namedexpr_test
734def is_simple_decorator_trailer(node: LN, last: bool = False) -> bool:
735 """Return True iff `node` is a trailer valid in a simple decorator"""
736 return node.type == syms.trailer and (
737 (
738 len(node.children) == 2
739 and node.children[0].type == token.DOT
740 and node.children[1].type == token.NAME
741 )
742 # last trailer can be an argument-less parentheses pair
743 or (
744 last
745 and len(node.children) == 2
746 and node.children[0].type == token.LPAR
747 and node.children[1].type == token.RPAR
748 )
749 # last trailer can be arguments
750 or (
751 last
752 and len(node.children) == 3
753 and node.children[0].type == token.LPAR
754 # and node.children[1].type == syms.argument
755 and node.children[2].type == token.RPAR
756 )
757 )
760def is_simple_decorator_expression(node: LN) -> bool:
761 """Return True iff `node` could be a 'dotted name' decorator
763 This function takes the node of the 'namedexpr_test' of the new decorator
764 grammar and test if it would be valid under the old decorator grammar.
766 The old grammar was: decorator: @ dotted_name [arguments] NEWLINE
767 The new grammar is : decorator: @ namedexpr_test NEWLINE
768 """
769 if node.type == token.NAME:
770 return True
771 if node.type == syms.power:
772 if node.children:
773 return (
774 node.children[0].type == token.NAME
775 and all(map(is_simple_decorator_trailer, node.children[1:-1]))
776 and (
777 len(node.children) < 2
778 or is_simple_decorator_trailer(node.children[-1], last=True)
779 )
780 )
781 return False
784def is_yield(node: LN) -> bool:
785 """Return True if `node` holds a `yield` or `yield from` expression."""
786 if node.type == syms.yield_expr:
787 return True
789 if is_name_token(node) and node.value == "yield":
790 return True
792 if node.type != syms.atom:
793 return False
795 if len(node.children) != 3:
796 return False
798 lpar, expr, rpar = node.children
799 if lpar.type == token.LPAR and rpar.type == token.RPAR:
800 return is_yield(expr)
802 return False
805def is_vararg(leaf: Leaf, within: set[NodeType]) -> bool:
806 """Return True if `leaf` is a star or double star in a vararg or kwarg.
808 If `within` includes VARARGS_PARENTS, this applies to function signatures.
809 If `within` includes UNPACKING_PARENTS, it applies to right hand-side
810 extended iterable unpacking (PEP 3132) and additional unpacking
811 generalizations (PEP 448).
812 """
813 if leaf.type not in VARARGS_SPECIALS or not leaf.parent:
814 return False
816 p = leaf.parent
817 if p.type == syms.star_expr:
818 # Star expressions are also used as assignment targets in extended
819 # iterable unpacking (PEP 3132). See what its parent is instead.
820 if not p.parent:
821 return False
823 p = p.parent
825 return p.type in within
828def is_fstring(node: Node) -> bool:
829 """Return True if the node is an f-string"""
830 return node.type == syms.fstring
833def fstring_tstring_to_string(node: Node) -> Leaf:
834 """Converts an fstring or tstring node back to a string node."""
835 string_without_prefix = str(node)[len(node.prefix) :]
836 string_leaf = Leaf(token.STRING, string_without_prefix, prefix=node.prefix)
837 string_leaf.lineno = node.get_lineno() or 0
838 return string_leaf
841def is_multiline_string(node: LN) -> bool:
842 """Return True if `leaf` is a multiline string that actually spans many lines."""
843 if isinstance(node, Node) and is_fstring(node):
844 leaf = fstring_tstring_to_string(node)
845 elif isinstance(node, Leaf):
846 leaf = node
847 else:
848 return False
850 return has_triple_quotes(leaf.value) and "\n" in leaf.value
853def is_parent_function_or_class(node: Node) -> bool:
854 assert node.type in {syms.suite, syms.simple_stmt}
855 assert node.parent is not None
856 # Note this works for suites / simple_stmts in async def as well
857 return node.parent.type in {syms.funcdef, syms.classdef}
860def is_stub_suite(node: Node) -> bool:
861 """Return True if `node` is a suite with a stub body."""
862 if node.parent is not None and not is_parent_function_or_class(node):
863 return False
865 # If there is a comment, we want to keep it.
866 if node.prefix.strip():
867 return False
869 if (
870 len(node.children) != 4
871 or node.children[0].type != token.NEWLINE
872 or node.children[1].type != token.INDENT
873 or node.children[3].type != token.DEDENT
874 ):
875 return False
877 if node.children[3].prefix.strip():
878 return False
880 return is_stub_body(node.children[2])
883def is_stub_body(node: LN) -> bool:
884 """Return True if `node` is a simple statement containing an ellipsis."""
885 if not isinstance(node, Node) or node.type != syms.simple_stmt:
886 return False
888 if len(node.children) != 2:
889 return False
891 child = node.children[0]
892 return (
893 not child.prefix.strip()
894 and child.type == syms.atom
895 and len(child.children) == 3
896 and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
897 )
900def is_atom_with_invisible_parens(node: LN) -> bool:
901 """Given a `LN`, determines whether it's an atom `node` with invisible
902 parens. Useful in dedupe-ing and normalizing parens.
903 """
904 if isinstance(node, Leaf) or node.type != syms.atom:
905 return False
907 first, last = node.children[0], node.children[-1]
908 return (
909 isinstance(first, Leaf)
910 and first.type == token.LPAR
911 and first.value == ""
912 and isinstance(last, Leaf)
913 and last.type == token.RPAR
914 and last.value == ""
915 )
918def is_empty_par(leaf: Leaf) -> bool:
919 return is_empty_lpar(leaf) or is_empty_rpar(leaf)
922def is_empty_lpar(leaf: Leaf) -> bool:
923 return leaf.type == token.LPAR and leaf.value == ""
926def is_empty_rpar(leaf: Leaf) -> bool:
927 return leaf.type == token.RPAR and leaf.value == ""
930def is_import(leaf: Leaf) -> bool:
931 """Return True if the given leaf starts an import statement."""
932 p = leaf.parent
933 t = leaf.type
934 v = leaf.value
935 return bool(
936 (t == token.LAZY and p and p.type == syms.lazy_import)
937 or (
938 t == token.NAME
939 and (
940 (v == "import" and p and p.type == syms.import_name)
941 or (v == "from" and p and p.type == syms.import_from)
942 )
943 )
944 )
947def is_with_or_async_with_stmt(leaf: Leaf) -> bool:
948 """Return True if the given leaf starts a with or async with statement."""
949 return bool(
950 leaf.type == token.NAME
951 and leaf.value == "with"
952 and leaf.parent
953 and leaf.parent.type == syms.with_stmt
954 ) or bool(
955 leaf.type == token.ASYNC
956 and leaf.next_sibling
957 and leaf.next_sibling.type == syms.with_stmt
958 )
961def is_async_stmt_or_funcdef(leaf: Leaf) -> bool:
962 """Return True if the given leaf starts an async def/for/with statement.
964 Note that `async def` can be either an `async_stmt` or `async_funcdef`,
965 the latter is used when it has decorators.
966 """
967 return bool(
968 leaf.type == token.ASYNC
969 and leaf.parent
970 and leaf.parent.type in {syms.async_stmt, syms.async_funcdef}
971 )
974def is_type_comment(leaf: Leaf, mode: Mode) -> bool:
975 """Return True if the given leaf is a type comment. This function should only
976 be used for general type comments (excluding ignore annotations, which should
977 use `is_type_ignore_comment`). Note that general type comments are no longer
978 used in modern version of Python, this function may be deprecated in the future."""
979 t = leaf.type
980 v = leaf.value
981 return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_comment_string(v, mode)
984def is_type_comment_string(value: str, mode: Mode) -> bool:
985 return value.startswith("#") and value[1:].lstrip().startswith("type:")
988def is_type_ignore_comment(leaf: Leaf, mode: Mode) -> bool:
989 """Return True if the given leaf is a type comment with ignore annotation."""
990 t = leaf.type
991 v = leaf.value
992 return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_ignore_comment_string(
993 v, mode
994 )
997def is_type_ignore_comment_string(value: str, mode: Mode) -> bool:
998 """Return True if the given string match with type comment with
999 ignore annotation."""
1000 return is_type_comment_string(value, mode) and value.split(":", 1)[
1001 1
1002 ].lstrip().startswith("ignore")
1005def wrap_in_parentheses(
1006 parent: Node, child: LN, *, visible: bool = True, index: int | None = None
1007) -> None:
1008 """Wrap `child` in parentheses.
1010 This replaces `child` with an atom holding the parentheses and the old
1011 child. That requires moving the prefix.
1013 If `visible` is False, the leaves will be valueless (and thus invisible).
1015 When the caller already knows `child`'s position in `parent.children` it
1016 can pass `index` so the child is swapped in place. The default path locates
1017 `child` with `Base.remove`, which scans and rewrites the whole child list;
1018 that is O(len(parent.children)) per call and turns quadratic when a caller
1019 wraps every child of a large node (e.g. each value of a big dict literal).
1020 """
1021 lpar = Leaf(token.LPAR, "(" if visible else "")
1022 rpar = Leaf(token.RPAR, ")" if visible else "")
1023 prefix = child.prefix
1024 child.prefix = ""
1025 if index is None:
1026 index = child.remove() or 0
1027 new_child = Node(syms.atom, [lpar, child, rpar])
1028 new_child.prefix = prefix
1029 parent.insert_child(index, new_child)
1030 else:
1031 # Detach the child pointer first so the atom can adopt it, then swap it
1032 # in place. set_child resets the old child's parent, so reattach it to
1033 # the new atom afterwards.
1034 child.parent = None
1035 new_child = Node(syms.atom, [lpar, child, rpar])
1036 new_child.prefix = prefix
1037 parent.set_child(index, new_child)
1038 child.parent = new_child
1041def unwrap_singleton_parenthesis(node: LN) -> LN | None:
1042 """Returns `wrapped` if `node` is of the shape ( wrapped ).
1044 Parenthesis can be optional. Returns None otherwise"""
1045 if len(node.children) != 3:
1046 return None
1048 lpar, wrapped, rpar = node.children
1049 if not (lpar.type == token.LPAR and rpar.type == token.RPAR):
1050 return None
1052 return wrapped
1055def ensure_visible(leaf: Leaf) -> None:
1056 """Make sure parentheses are visible.
1058 They could be invisible as part of some statements (see
1059 :func:`normalize_invisible_parens` and :func:`visit_import_from`).
1060 """
1061 if leaf.type == token.LPAR:
1062 leaf.value = "("
1063 elif leaf.type == token.RPAR:
1064 leaf.value = ")"
1067def is_name_token(nl: NL) -> TypeGuard[Leaf]:
1068 return nl.type == token.NAME
1071def is_lpar_token(nl: NL) -> TypeGuard[Leaf]:
1072 return nl.type == token.LPAR
1075def is_rpar_token(nl: NL) -> TypeGuard[Leaf]:
1076 return nl.type == token.RPAR
1079def is_number_token(nl: NL) -> TypeGuard[Leaf]:
1080 return nl.type == token.NUMBER
1083def get_annotation_type(leaf: Leaf) -> Literal["return", "param", None]:
1084 """Returns the type of annotation this leaf is part of, if any."""
1085 ancestor = leaf.parent
1086 while ancestor is not None:
1087 if ancestor.prev_sibling and ancestor.prev_sibling.type == token.RARROW:
1088 return "return"
1089 if ancestor.parent and ancestor.parent.type == syms.tname:
1090 return "param"
1091 ancestor = ancestor.parent
1092 return None
1095def is_part_of_annotation(leaf: Leaf) -> bool:
1096 """Returns whether this leaf is part of a type annotation."""
1097 assert leaf.parent is not None
1098 return get_annotation_type(leaf) is not None
1101def first_leaf(node: LN) -> Leaf | None:
1102 """Returns the first leaf of the ancestor node."""
1103 if isinstance(node, Leaf):
1104 return node
1105 elif not node.children:
1106 return None
1107 else:
1108 return first_leaf(node.children[0])
1111def last_leaf(node: LN) -> Leaf | None:
1112 """Returns the last leaf of the ancestor node."""
1113 if isinstance(node, Leaf):
1114 return node
1115 elif not node.children:
1116 return None
1117 else:
1118 return last_leaf(node.children[-1])
1121def furthest_ancestor_with_last_leaf(leaf: Leaf) -> LN:
1122 """Returns the furthest ancestor that has this leaf node as the last leaf."""
1123 node: LN = leaf
1124 while node.parent and node.parent.children and node is node.parent.children[-1]:
1125 node = node.parent
1126 return node
1129def has_sibling_with_type(node: LN, type: int) -> bool:
1130 # Check previous siblings
1131 sibling = node.prev_sibling
1132 while sibling is not None:
1133 if sibling.type == type:
1134 return True
1135 sibling = sibling.prev_sibling
1137 # Check next siblings
1138 sibling = node.next_sibling
1139 while sibling is not None:
1140 if sibling.type == type:
1141 return True
1142 sibling = sibling.next_sibling
1144 return False