]> git.madduck.net Git - etc/vim.git/blob - src/black/nodes.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:

45070909df466c2fce7746788e6ac762c5a091a7
[etc/vim.git] / src / black / nodes.py
1 """
2 blib2to3 Node/Leaf transformation-related utility functions.
3 """
4
5 import sys
6 from typing import Generic, Iterator, List, Optional, Set, Tuple, TypeVar, Union
7
8 if sys.version_info >= (3, 8):
9     from typing import Final
10 else:
11     from typing_extensions import Final
12 if sys.version_info >= (3, 10):
13     from typing import TypeGuard
14 else:
15     from typing_extensions import TypeGuard
16
17 from mypy_extensions import mypyc_attr
18
19 from black.cache import CACHE_DIR
20 from black.strings import has_triple_quotes
21 from blib2to3 import pygram
22 from blib2to3.pgen2 import token
23 from blib2to3.pytree import NL, Leaf, Node, type_repr
24
25 pygram.initialize(CACHE_DIR)
26 syms: Final = pygram.python_symbols
27
28
29 # types
30 T = TypeVar("T")
31 LN = Union[Leaf, Node]
32 LeafID = int
33 NodeType = int
34
35
36 WHITESPACE: Final = {token.DEDENT, token.INDENT, token.NEWLINE}
37 STATEMENT: Final = {
38     syms.if_stmt,
39     syms.while_stmt,
40     syms.for_stmt,
41     syms.try_stmt,
42     syms.except_clause,
43     syms.with_stmt,
44     syms.funcdef,
45     syms.classdef,
46     syms.match_stmt,
47     syms.case_block,
48 }
49 STANDALONE_COMMENT: Final = 153
50 token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"
51 LOGIC_OPERATORS: Final = {"and", "or"}
52 COMPARATORS: Final = {
53     token.LESS,
54     token.GREATER,
55     token.EQEQUAL,
56     token.NOTEQUAL,
57     token.LESSEQUAL,
58     token.GREATEREQUAL,
59 }
60 MATH_OPERATORS: Final = {
61     token.VBAR,
62     token.CIRCUMFLEX,
63     token.AMPER,
64     token.LEFTSHIFT,
65     token.RIGHTSHIFT,
66     token.PLUS,
67     token.MINUS,
68     token.STAR,
69     token.SLASH,
70     token.DOUBLESLASH,
71     token.PERCENT,
72     token.AT,
73     token.TILDE,
74     token.DOUBLESTAR,
75 }
76 STARS: Final = {token.STAR, token.DOUBLESTAR}
77 VARARGS_SPECIALS: Final = STARS | {token.SLASH}
78 VARARGS_PARENTS: Final = {
79     syms.arglist,
80     syms.argument,  # double star in arglist
81     syms.trailer,  # single argument to call
82     syms.typedargslist,
83     syms.varargslist,  # lambdas
84 }
85 UNPACKING_PARENTS: Final = {
86     syms.atom,  # single element of a list or set literal
87     syms.dictsetmaker,
88     syms.listmaker,
89     syms.testlist_gexp,
90     syms.testlist_star_expr,
91     syms.subject_expr,
92     syms.pattern,
93 }
94 TEST_DESCENDANTS: Final = {
95     syms.test,
96     syms.lambdef,
97     syms.or_test,
98     syms.and_test,
99     syms.not_test,
100     syms.comparison,
101     syms.star_expr,
102     syms.expr,
103     syms.xor_expr,
104     syms.and_expr,
105     syms.shift_expr,
106     syms.arith_expr,
107     syms.trailer,
108     syms.term,
109     syms.power,
110 }
111 TYPED_NAMES: Final = {syms.tname, syms.tname_star}
112 ASSIGNMENTS: Final = {
113     "=",
114     "+=",
115     "-=",
116     "*=",
117     "@=",
118     "/=",
119     "%=",
120     "&=",
121     "|=",
122     "^=",
123     "<<=",
124     ">>=",
125     "**=",
126     "//=",
127 }
128
129 IMPLICIT_TUPLE: Final = {syms.testlist, syms.testlist_star_expr, syms.exprlist}
130 BRACKET: Final = {
131     token.LPAR: token.RPAR,
132     token.LSQB: token.RSQB,
133     token.LBRACE: token.RBRACE,
134 }
135 OPENING_BRACKETS: Final = set(BRACKET.keys())
136 CLOSING_BRACKETS: Final = set(BRACKET.values())
137 BRACKETS: Final = OPENING_BRACKETS | CLOSING_BRACKETS
138 ALWAYS_NO_SPACE: Final = CLOSING_BRACKETS | {token.COMMA, STANDALONE_COMMENT}
139
140 RARROW = 55
141
142
143 @mypyc_attr(allow_interpreted_subclasses=True)
144 class Visitor(Generic[T]):
145     """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""
146
147     def visit(self, node: LN) -> Iterator[T]:
148         """Main method to visit `node` and its children.
149
150         It tries to find a `visit_*()` method for the given `node.type`, like
151         `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.
152         If no dedicated `visit_*()` method is found, chooses `visit_default()`
153         instead.
154
155         Then yields objects of type `T` from the selected visitor.
156         """
157         if node.type < 256:
158             name = token.tok_name[node.type]
159         else:
160             name = str(type_repr(node.type))
161         # We explicitly branch on whether a visitor exists (instead of
162         # using self.visit_default as the default arg to getattr) in order
163         # to save needing to create a bound method object and so mypyc can
164         # generate a native call to visit_default.
165         visitf = getattr(self, f"visit_{name}", None)
166         if visitf:
167             yield from visitf(node)
168         else:
169             yield from self.visit_default(node)
170
171     def visit_default(self, node: LN) -> Iterator[T]:
172         """Default `visit_*()` implementation. Recurses to children of `node`."""
173         if isinstance(node, Node):
174             for child in node.children:
175                 yield from self.visit(child)
176
177
178 def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str:  # noqa: C901
179     """Return whitespace prefix if needed for the given `leaf`.
180
181     `complex_subscript` signals whether the given leaf is part of a subscription
182     which has non-trivial arguments, like arithmetic expressions or function calls.
183     """
184     NO: Final[str] = ""
185     SPACE: Final[str] = " "
186     DOUBLESPACE: Final[str] = "  "
187     t = leaf.type
188     p = leaf.parent
189     v = leaf.value
190     if t in ALWAYS_NO_SPACE:
191         return NO
192
193     if t == token.COMMENT:
194         return DOUBLESPACE
195
196     assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
197     if t == token.COLON and p.type not in {
198         syms.subscript,
199         syms.subscriptlist,
200         syms.sliceop,
201     }:
202         return NO
203
204     prev = leaf.prev_sibling
205     if not prev:
206         prevp = preceding_leaf(p)
207         if not prevp or prevp.type in OPENING_BRACKETS:
208             return NO
209
210         if t == token.COLON:
211             if prevp.type == token.COLON:
212                 return NO
213
214             elif prevp.type != token.COMMA and not complex_subscript:
215                 return NO
216
217             return SPACE
218
219         if prevp.type == token.EQUAL:
220             if prevp.parent:
221                 if prevp.parent.type in {
222                     syms.arglist,
223                     syms.argument,
224                     syms.parameters,
225                     syms.varargslist,
226                 }:
227                     return NO
228
229                 elif prevp.parent.type == syms.typedargslist:
230                     # A bit hacky: if the equal sign has whitespace, it means we
231                     # previously found it's a typed argument.  So, we're using
232                     # that, too.
233                     return prevp.prefix
234
235         elif (
236             prevp.type == token.STAR
237             and parent_type(prevp) == syms.star_expr
238             and parent_type(prevp.parent) == syms.subscriptlist
239         ):
240             # No space between typevar tuples.
241             return NO
242
243         elif prevp.type in VARARGS_SPECIALS:
244             if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
245                 return NO
246
247         elif prevp.type == token.COLON:
248             if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:
249                 return SPACE if complex_subscript else NO
250
251         elif (
252             prevp.parent
253             and prevp.parent.type == syms.factor
254             and prevp.type in MATH_OPERATORS
255         ):
256             return NO
257
258         elif prevp.type == token.AT and p.parent and p.parent.type == syms.decorator:
259             # no space in decorators
260             return NO
261
262     elif prev.type in OPENING_BRACKETS:
263         return NO
264
265     if p.type in {syms.parameters, syms.arglist}:
266         # untyped function signatures or calls
267         if not prev or prev.type != token.COMMA:
268             return NO
269
270     elif p.type == syms.varargslist:
271         # lambdas
272         if prev and prev.type != token.COMMA:
273             return NO
274
275     elif p.type == syms.typedargslist:
276         # typed function signatures
277         if not prev:
278             return NO
279
280         if t == token.EQUAL:
281             if prev.type not in TYPED_NAMES:
282                 return NO
283
284         elif prev.type == token.EQUAL:
285             # A bit hacky: if the equal sign has whitespace, it means we
286             # previously found it's a typed argument.  So, we're using that, too.
287             return prev.prefix
288
289         elif prev.type != token.COMMA:
290             return NO
291
292     elif p.type in TYPED_NAMES:
293         # type names
294         if not prev:
295             prevp = preceding_leaf(p)
296             if not prevp or prevp.type != token.COMMA:
297                 return NO
298
299     elif p.type == syms.trailer:
300         # attributes and calls
301         if t == token.LPAR or t == token.RPAR:
302             return NO
303
304         if not prev:
305             if t == token.DOT or t == token.LSQB:
306                 return NO
307
308         elif prev.type != token.COMMA:
309             return NO
310
311     elif p.type == syms.argument:
312         # single argument
313         if t == token.EQUAL:
314             return NO
315
316         if not prev:
317             prevp = preceding_leaf(p)
318             if not prevp or prevp.type == token.LPAR:
319                 return NO
320
321         elif prev.type in {token.EQUAL} | VARARGS_SPECIALS:
322             return NO
323
324     elif p.type == syms.decorator:
325         # decorators
326         return NO
327
328     elif p.type == syms.dotted_name:
329         if prev:
330             return NO
331
332         prevp = preceding_leaf(p)
333         if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
334             return NO
335
336     elif p.type == syms.classdef:
337         if t == token.LPAR:
338             return NO
339
340         if prev and prev.type == token.LPAR:
341             return NO
342
343     elif p.type in {syms.subscript, syms.sliceop}:
344         # indexing
345         if not prev:
346             assert p.parent is not None, "subscripts are always parented"
347             if p.parent.type == syms.subscriptlist:
348                 return SPACE
349
350             return NO
351
352         elif not complex_subscript:
353             return NO
354
355     elif p.type == syms.atom:
356         if prev and t == token.DOT:
357             # dots, but not the first one.
358             return NO
359
360     elif p.type == syms.dictsetmaker:
361         # dict unpacking
362         if prev and prev.type == token.DOUBLESTAR:
363             return NO
364
365     elif p.type in {syms.factor, syms.star_expr}:
366         # unary ops
367         if not prev:
368             prevp = preceding_leaf(p)
369             if not prevp or prevp.type in OPENING_BRACKETS:
370                 return NO
371
372             prevp_parent = prevp.parent
373             assert prevp_parent is not None
374             if prevp.type == token.COLON and prevp_parent.type in {
375                 syms.subscript,
376                 syms.sliceop,
377             }:
378                 return NO
379
380             elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
381                 return NO
382
383         elif t in {token.NAME, token.NUMBER, token.STRING}:
384             return NO
385
386     elif p.type == syms.import_from:
387         if t == token.DOT:
388             if prev and prev.type == token.DOT:
389                 return NO
390
391         elif t == token.NAME:
392             if v == "import":
393                 return SPACE
394
395             if prev and prev.type == token.DOT:
396                 return NO
397
398     elif p.type == syms.sliceop:
399         return NO
400
401     elif p.type == syms.except_clause:
402         if t == token.STAR:
403             return NO
404
405     return SPACE
406
407
408 def preceding_leaf(node: Optional[LN]) -> Optional[Leaf]:
409     """Return the first leaf that precedes `node`, if any."""
410     while node:
411         res = node.prev_sibling
412         if res:
413             if isinstance(res, Leaf):
414                 return res
415
416             try:
417                 return list(res.leaves())[-1]
418
419             except IndexError:
420                 return None
421
422         node = node.parent
423     return None
424
425
426 def prev_siblings_are(node: Optional[LN], tokens: List[Optional[NodeType]]) -> bool:
427     """Return if the `node` and its previous siblings match types against the provided
428     list of tokens; the provided `node`has its type matched against the last element in
429     the list.  `None` can be used as the first element to declare that the start of the
430     list is anchored at the start of its parent's children."""
431     if not tokens:
432         return True
433     if tokens[-1] is None:
434         return node is None
435     if not node:
436         return False
437     if node.type != tokens[-1]:
438         return False
439     return prev_siblings_are(node.prev_sibling, tokens[:-1])
440
441
442 def parent_type(node: Optional[LN]) -> Optional[NodeType]:
443     """
444     Returns:
445         @node.parent.type, if @node is not None and has a parent.
446             OR
447         None, otherwise.
448     """
449     if node is None or node.parent is None:
450         return None
451
452     return node.parent.type
453
454
455 def child_towards(ancestor: Node, descendant: LN) -> Optional[LN]:
456     """Return the child of `ancestor` that contains `descendant`."""
457     node: Optional[LN] = descendant
458     while node and node.parent != ancestor:
459         node = node.parent
460     return node
461
462
463 def replace_child(old_child: LN, new_child: LN) -> None:
464     """
465     Side Effects:
466         * If @old_child.parent is set, replace @old_child with @new_child in
467         @old_child's underlying Node structure.
468             OR
469         * Otherwise, this function does nothing.
470     """
471     parent = old_child.parent
472     if not parent:
473         return
474
475     child_idx = old_child.remove()
476     if child_idx is not None:
477         parent.insert_child(child_idx, new_child)
478
479
480 def container_of(leaf: Leaf) -> LN:
481     """Return `leaf` or one of its ancestors that is the topmost container of it.
482
483     By "container" we mean a node where `leaf` is the very first child.
484     """
485     same_prefix = leaf.prefix
486     container: LN = leaf
487     while container:
488         parent = container.parent
489         if parent is None:
490             break
491
492         if parent.children[0].prefix != same_prefix:
493             break
494
495         if parent.type == syms.file_input:
496             break
497
498         if parent.prev_sibling is not None and parent.prev_sibling.type in BRACKETS:
499             break
500
501         container = parent
502     return container
503
504
505 def first_leaf_of(node: LN) -> Optional[Leaf]:
506     """Returns the first leaf of the node tree."""
507     if isinstance(node, Leaf):
508         return node
509     if node.children:
510         return first_leaf_of(node.children[0])
511     else:
512         return None
513
514
515 def is_arith_like(node: LN) -> bool:
516     """Whether node is an arithmetic or a binary arithmetic expression"""
517     return node.type in {
518         syms.arith_expr,
519         syms.shift_expr,
520         syms.xor_expr,
521         syms.and_expr,
522     }
523
524
525 def is_docstring(leaf: Leaf) -> bool:
526     if prev_siblings_are(
527         leaf.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
528     ):
529         return True
530
531     # Multiline docstring on the same line as the `def`.
532     if prev_siblings_are(leaf.parent, [syms.parameters, token.COLON, syms.simple_stmt]):
533         # `syms.parameters` is only used in funcdefs and async_funcdefs in the Python
534         # grammar. We're safe to return True without further checks.
535         return True
536
537     return False
538
539
540 def is_empty_tuple(node: LN) -> bool:
541     """Return True if `node` holds an empty tuple."""
542     return (
543         node.type == syms.atom
544         and len(node.children) == 2
545         and node.children[0].type == token.LPAR
546         and node.children[1].type == token.RPAR
547     )
548
549
550 def is_one_tuple(node: LN) -> bool:
551     """Return True if `node` holds a tuple with one element, with or without parens."""
552     if node.type == syms.atom:
553         gexp = unwrap_singleton_parenthesis(node)
554         if gexp is None or gexp.type != syms.testlist_gexp:
555             return False
556
557         return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA
558
559     return (
560         node.type in IMPLICIT_TUPLE
561         and len(node.children) == 2
562         and node.children[1].type == token.COMMA
563     )
564
565
566 def is_tuple_containing_walrus(node: LN) -> bool:
567     """Return True if `node` holds a tuple that contains a walrus operator."""
568     if node.type != syms.atom:
569         return False
570     gexp = unwrap_singleton_parenthesis(node)
571     if gexp is None or gexp.type != syms.testlist_gexp:
572         return False
573
574     return any(child.type == syms.namedexpr_test for child in gexp.children)
575
576
577 def is_one_sequence_between(
578     opening: Leaf,
579     closing: Leaf,
580     leaves: List[Leaf],
581     brackets: Tuple[int, int] = (token.LPAR, token.RPAR),
582 ) -> bool:
583     """Return True if content between `opening` and `closing` is a one-sequence."""
584     if (opening.type, closing.type) != brackets:
585         return False
586
587     depth = closing.bracket_depth + 1
588     for _opening_index, leaf in enumerate(leaves):
589         if leaf is opening:
590             break
591
592     else:
593         raise LookupError("Opening paren not found in `leaves`")
594
595     commas = 0
596     _opening_index += 1
597     for leaf in leaves[_opening_index:]:
598         if leaf is closing:
599             break
600
601         bracket_depth = leaf.bracket_depth
602         if bracket_depth == depth and leaf.type == token.COMMA:
603             commas += 1
604             if leaf.parent and leaf.parent.type in {
605                 syms.arglist,
606                 syms.typedargslist,
607             }:
608                 commas += 1
609                 break
610
611     return commas < 2
612
613
614 def is_walrus_assignment(node: LN) -> bool:
615     """Return True iff `node` is of the shape ( test := test )"""
616     inner = unwrap_singleton_parenthesis(node)
617     return inner is not None and inner.type == syms.namedexpr_test
618
619
620 def is_simple_decorator_trailer(node: LN, last: bool = False) -> bool:
621     """Return True iff `node` is a trailer valid in a simple decorator"""
622     return node.type == syms.trailer and (
623         (
624             len(node.children) == 2
625             and node.children[0].type == token.DOT
626             and node.children[1].type == token.NAME
627         )
628         # last trailer can be an argument-less parentheses pair
629         or (
630             last
631             and len(node.children) == 2
632             and node.children[0].type == token.LPAR
633             and node.children[1].type == token.RPAR
634         )
635         # last trailer can be arguments
636         or (
637             last
638             and len(node.children) == 3
639             and node.children[0].type == token.LPAR
640             # and node.children[1].type == syms.argument
641             and node.children[2].type == token.RPAR
642         )
643     )
644
645
646 def is_simple_decorator_expression(node: LN) -> bool:
647     """Return True iff `node` could be a 'dotted name' decorator
648
649     This function takes the node of the 'namedexpr_test' of the new decorator
650     grammar and test if it would be valid under the old decorator grammar.
651
652     The old grammar was: decorator: @ dotted_name [arguments] NEWLINE
653     The new grammar is : decorator: @ namedexpr_test NEWLINE
654     """
655     if node.type == token.NAME:
656         return True
657     if node.type == syms.power:
658         if node.children:
659             return (
660                 node.children[0].type == token.NAME
661                 and all(map(is_simple_decorator_trailer, node.children[1:-1]))
662                 and (
663                     len(node.children) < 2
664                     or is_simple_decorator_trailer(node.children[-1], last=True)
665                 )
666             )
667     return False
668
669
670 def is_yield(node: LN) -> bool:
671     """Return True if `node` holds a `yield` or `yield from` expression."""
672     if node.type == syms.yield_expr:
673         return True
674
675     if is_name_token(node) and node.value == "yield":
676         return True
677
678     if node.type != syms.atom:
679         return False
680
681     if len(node.children) != 3:
682         return False
683
684     lpar, expr, rpar = node.children
685     if lpar.type == token.LPAR and rpar.type == token.RPAR:
686         return is_yield(expr)
687
688     return False
689
690
691 def is_vararg(leaf: Leaf, within: Set[NodeType]) -> bool:
692     """Return True if `leaf` is a star or double star in a vararg or kwarg.
693
694     If `within` includes VARARGS_PARENTS, this applies to function signatures.
695     If `within` includes UNPACKING_PARENTS, it applies to right hand-side
696     extended iterable unpacking (PEP 3132) and additional unpacking
697     generalizations (PEP 448).
698     """
699     if leaf.type not in VARARGS_SPECIALS or not leaf.parent:
700         return False
701
702     p = leaf.parent
703     if p.type == syms.star_expr:
704         # Star expressions are also used as assignment targets in extended
705         # iterable unpacking (PEP 3132).  See what its parent is instead.
706         if not p.parent:
707             return False
708
709         p = p.parent
710
711     return p.type in within
712
713
714 def is_multiline_string(leaf: Leaf) -> bool:
715     """Return True if `leaf` is a multiline string that actually spans many lines."""
716     return has_triple_quotes(leaf.value) and "\n" in leaf.value
717
718
719 def is_stub_suite(node: Node) -> bool:
720     """Return True if `node` is a suite with a stub body."""
721     if (
722         len(node.children) != 4
723         or node.children[0].type != token.NEWLINE
724         or node.children[1].type != token.INDENT
725         or node.children[3].type != token.DEDENT
726     ):
727         return False
728
729     return is_stub_body(node.children[2])
730
731
732 def is_stub_body(node: LN) -> bool:
733     """Return True if `node` is a simple statement containing an ellipsis."""
734     if not isinstance(node, Node) or node.type != syms.simple_stmt:
735         return False
736
737     if len(node.children) != 2:
738         return False
739
740     child = node.children[0]
741     return (
742         child.type == syms.atom
743         and len(child.children) == 3
744         and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
745     )
746
747
748 def is_atom_with_invisible_parens(node: LN) -> bool:
749     """Given a `LN`, determines whether it's an atom `node` with invisible
750     parens. Useful in dedupe-ing and normalizing parens.
751     """
752     if isinstance(node, Leaf) or node.type != syms.atom:
753         return False
754
755     first, last = node.children[0], node.children[-1]
756     return (
757         isinstance(first, Leaf)
758         and first.type == token.LPAR
759         and first.value == ""
760         and isinstance(last, Leaf)
761         and last.type == token.RPAR
762         and last.value == ""
763     )
764
765
766 def is_empty_par(leaf: Leaf) -> bool:
767     return is_empty_lpar(leaf) or is_empty_rpar(leaf)
768
769
770 def is_empty_lpar(leaf: Leaf) -> bool:
771     return leaf.type == token.LPAR and leaf.value == ""
772
773
774 def is_empty_rpar(leaf: Leaf) -> bool:
775     return leaf.type == token.RPAR and leaf.value == ""
776
777
778 def is_import(leaf: Leaf) -> bool:
779     """Return True if the given leaf starts an import statement."""
780     p = leaf.parent
781     t = leaf.type
782     v = leaf.value
783     return bool(
784         t == token.NAME
785         and (
786             (v == "import" and p and p.type == syms.import_name)
787             or (v == "from" and p and p.type == syms.import_from)
788         )
789     )
790
791
792 def is_with_or_async_with_stmt(leaf: Leaf) -> bool:
793     """Return True if the given leaf starts a with or async with statement."""
794     return bool(
795         leaf.type == token.NAME
796         and leaf.value == "with"
797         and leaf.parent
798         and leaf.parent.type == syms.with_stmt
799     ) or bool(
800         leaf.type == token.ASYNC
801         and leaf.next_sibling
802         and leaf.next_sibling.type == syms.with_stmt
803     )
804
805
806 def is_async_stmt_or_funcdef(leaf: Leaf) -> bool:
807     """Return True if the given leaf starts an async def/for/with statement.
808
809     Note that `async def` can be either an `async_stmt` or `async_funcdef`,
810     the latter is used when it has decorators.
811     """
812     return bool(
813         leaf.type == token.ASYNC
814         and leaf.parent
815         and leaf.parent.type in {syms.async_stmt, syms.async_funcdef}
816     )
817
818
819 def is_type_comment(leaf: Leaf, suffix: str = "") -> bool:
820     """Return True if the given leaf is a special comment.
821     Only returns true for type comments for now."""
822     t = leaf.type
823     v = leaf.value
824     return t in {token.COMMENT, STANDALONE_COMMENT} and v.startswith("# type:" + suffix)
825
826
827 def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:
828     """Wrap `child` in parentheses.
829
830     This replaces `child` with an atom holding the parentheses and the old
831     child.  That requires moving the prefix.
832
833     If `visible` is False, the leaves will be valueless (and thus invisible).
834     """
835     lpar = Leaf(token.LPAR, "(" if visible else "")
836     rpar = Leaf(token.RPAR, ")" if visible else "")
837     prefix = child.prefix
838     child.prefix = ""
839     index = child.remove() or 0
840     new_child = Node(syms.atom, [lpar, child, rpar])
841     new_child.prefix = prefix
842     parent.insert_child(index, new_child)
843
844
845 def unwrap_singleton_parenthesis(node: LN) -> Optional[LN]:
846     """Returns `wrapped` if `node` is of the shape ( wrapped ).
847
848     Parenthesis can be optional. Returns None otherwise"""
849     if len(node.children) != 3:
850         return None
851
852     lpar, wrapped, rpar = node.children
853     if not (lpar.type == token.LPAR and rpar.type == token.RPAR):
854         return None
855
856     return wrapped
857
858
859 def ensure_visible(leaf: Leaf) -> None:
860     """Make sure parentheses are visible.
861
862     They could be invisible as part of some statements (see
863     :func:`normalize_invisible_parens` and :func:`visit_import_from`).
864     """
865     if leaf.type == token.LPAR:
866         leaf.value = "("
867     elif leaf.type == token.RPAR:
868         leaf.value = ")"
869
870
871 def is_name_token(nl: NL) -> TypeGuard[Leaf]:
872     return nl.type == token.NAME
873
874
875 def is_lpar_token(nl: NL) -> TypeGuard[Leaf]:
876     return nl.type == token.LPAR
877
878
879 def is_rpar_token(nl: NL) -> TypeGuard[Leaf]:
880     return nl.type == token.RPAR
881
882
883 def is_string_token(nl: NL) -> TypeGuard[Leaf]:
884     return nl.type == token.STRING
885
886
887 def is_number_token(nl: NL) -> TypeGuard[Leaf]:
888     return nl.type == token.NUMBER
889
890
891 def is_part_of_annotation(leaf: Leaf) -> bool:
892     """Returns whether this leaf is part of type annotations."""
893     ancestor = leaf.parent
894     while ancestor is not None:
895         if ancestor.prev_sibling and ancestor.prev_sibling.type == token.RARROW:
896             return True
897         if ancestor.parent and ancestor.parent.type == syms.tname:
898             return True
899         ancestor = ancestor.parent
900     return False