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 = ""
 185     SPACE: Final = " "
 186     DOUBLESPACE: Final = "  "
 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_column(node: Node) -> Optional[int]:
 506     """Returns the column of the first leaf child of a node."""
 507     for child in node.children:
 508         if isinstance(child, Leaf):
 509             return child.column
 510     return None
 511
 512
 513 def is_arith_like(node: LN) -> bool:
 514     """Whether node is an arithmetic or a binary arithmetic expression"""
 515     return node.type in {
 516         syms.arith_expr,
 517         syms.shift_expr,
 518         syms.xor_expr,
 519         syms.and_expr,
 520     }
 521
 522
 523 def is_docstring(leaf: Leaf) -> bool:
 524     if prev_siblings_are(
 525         leaf.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
 526     ):
 527         return True
 528
 529     # Multiline docstring on the same line as the `def`.
 530     if prev_siblings_are(leaf.parent, [syms.parameters, token.COLON, syms.simple_stmt]):
 531         # `syms.parameters` is only used in funcdefs and async_funcdefs in the Python
 532         # grammar. We're safe to return True without further checks.
 533         return True
 534
 535     return False
 536
 537
 538 def is_empty_tuple(node: LN) -> bool:
 539     """Return True if `node` holds an empty tuple."""
 540     return (
 541         node.type == syms.atom
 542         and len(node.children) == 2
 543         and node.children[0].type == token.LPAR
 544         and node.children[1].type == token.RPAR
 545     )
 546
 547
 548 def is_one_tuple(node: LN) -> bool:
 549     """Return True if `node` holds a tuple with one element, with or without parens."""
 550     if node.type == syms.atom:
 551         gexp = unwrap_singleton_parenthesis(node)
 552         if gexp is None or gexp.type != syms.testlist_gexp:
 553             return False
 554
 555         return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA
 556
 557     return (
 558         node.type in IMPLICIT_TUPLE
 559         and len(node.children) == 2
 560         and node.children[1].type == token.COMMA
 561     )
 562
 563
 564 def is_one_sequence_between(
 565     opening: Leaf,
 566     closing: Leaf,
 567     leaves: List[Leaf],
 568     brackets: Tuple[int, int] = (token.LPAR, token.RPAR),
 569 ) -> bool:
 570     """Return True if content between `opening` and `closing` is a one-sequence."""
 571     if (opening.type, closing.type) != brackets:
 572         return False
 573
 574     depth = closing.bracket_depth + 1
 575     for _opening_index, leaf in enumerate(leaves):
 576         if leaf is opening:
 577             break
 578
 579     else:
 580         raise LookupError("Opening paren not found in `leaves`")
 581
 582     commas = 0
 583     _opening_index += 1
 584     for leaf in leaves[_opening_index:]:
 585         if leaf is closing:
 586             break
 587
 588         bracket_depth = leaf.bracket_depth
 589         if bracket_depth == depth and leaf.type == token.COMMA:
 590             commas += 1
 591             if leaf.parent and leaf.parent.type in {
 592                 syms.arglist,
 593                 syms.typedargslist,
 594             }:
 595                 commas += 1
 596                 break
 597
 598     return commas < 2
 599
 600
 601 def is_walrus_assignment(node: LN) -> bool:
 602     """Return True iff `node` is of the shape ( test := test )"""
 603     inner = unwrap_singleton_parenthesis(node)
 604     return inner is not None and inner.type == syms.namedexpr_test
 605
 606
 607 def is_simple_decorator_trailer(node: LN, last: bool = False) -> bool:
 608     """Return True iff `node` is a trailer valid in a simple decorator"""
 609     return node.type == syms.trailer and (
 610         (
 611             len(node.children) == 2
 612             and node.children[0].type == token.DOT
 613             and node.children[1].type == token.NAME
 614         )
 615         # last trailer can be an argument-less parentheses pair
 616         or (
 617             last
 618             and len(node.children) == 2
 619             and node.children[0].type == token.LPAR
 620             and node.children[1].type == token.RPAR
 621         )
 622         # last trailer can be arguments
 623         or (
 624             last
 625             and len(node.children) == 3
 626             and node.children[0].type == token.LPAR
 627             # and node.children[1].type == syms.argument
 628             and node.children[2].type == token.RPAR
 629         )
 630     )
 631
 632
 633 def is_simple_decorator_expression(node: LN) -> bool:
 634     """Return True iff `node` could be a 'dotted name' decorator
 635
 636     This function takes the node of the 'namedexpr_test' of the new decorator
 637     grammar and test if it would be valid under the old decorator grammar.
 638
 639     The old grammar was: decorator: @ dotted_name [arguments] NEWLINE
 640     The new grammar is : decorator: @ namedexpr_test NEWLINE
 641     """
 642     if node.type == token.NAME:
 643         return True
 644     if node.type == syms.power:
 645         if node.children:
 646             return (
 647                 node.children[0].type == token.NAME
 648                 and all(map(is_simple_decorator_trailer, node.children[1:-1]))
 649                 and (
 650                     len(node.children) < 2
 651                     or is_simple_decorator_trailer(node.children[-1], last=True)
 652                 )
 653             )
 654     return False
 655
 656
 657 def is_yield(node: LN) -> bool:
 658     """Return True if `node` holds a `yield` or `yield from` expression."""
 659     if node.type == syms.yield_expr:
 660         return True
 661
 662     if is_name_token(node) and node.value == "yield":
 663         return True
 664
 665     if node.type != syms.atom:
 666         return False
 667
 668     if len(node.children) != 3:
 669         return False
 670
 671     lpar, expr, rpar = node.children
 672     if lpar.type == token.LPAR and rpar.type == token.RPAR:
 673         return is_yield(expr)
 674
 675     return False
 676
 677
 678 def is_vararg(leaf: Leaf, within: Set[NodeType]) -> bool:
 679     """Return True if `leaf` is a star or double star in a vararg or kwarg.
 680
 681     If `within` includes VARARGS_PARENTS, this applies to function signatures.
 682     If `within` includes UNPACKING_PARENTS, it applies to right hand-side
 683     extended iterable unpacking (PEP 3132) and additional unpacking
 684     generalizations (PEP 448).
 685     """
 686     if leaf.type not in VARARGS_SPECIALS or not leaf.parent:
 687         return False
 688
 689     p = leaf.parent
 690     if p.type == syms.star_expr:
 691         # Star expressions are also used as assignment targets in extended
 692         # iterable unpacking (PEP 3132).  See what its parent is instead.
 693         if not p.parent:
 694             return False
 695
 696         p = p.parent
 697
 698     return p.type in within
 699
 700
 701 def is_multiline_string(leaf: Leaf) -> bool:
 702     """Return True if `leaf` is a multiline string that actually spans many lines."""
 703     return has_triple_quotes(leaf.value) and "\n" in leaf.value
 704
 705
 706 def is_stub_suite(node: Node) -> bool:
 707     """Return True if `node` is a suite with a stub body."""
 708     if (
 709         len(node.children) != 4
 710         or node.children[0].type != token.NEWLINE
 711         or node.children[1].type != token.INDENT
 712         or node.children[3].type != token.DEDENT
 713     ):
 714         return False
 715
 716     return is_stub_body(node.children[2])
 717
 718
 719 def is_stub_body(node: LN) -> bool:
 720     """Return True if `node` is a simple statement containing an ellipsis."""
 721     if not isinstance(node, Node) or node.type != syms.simple_stmt:
 722         return False
 723
 724     if len(node.children) != 2:
 725         return False
 726
 727     child = node.children[0]
 728     return (
 729         child.type == syms.atom
 730         and len(child.children) == 3
 731         and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
 732     )
 733
 734
 735 def is_atom_with_invisible_parens(node: LN) -> bool:
 736     """Given a `LN`, determines whether it's an atom `node` with invisible
 737     parens. Useful in dedupe-ing and normalizing parens.
 738     """
 739     if isinstance(node, Leaf) or node.type != syms.atom:
 740         return False
 741
 742     first, last = node.children[0], node.children[-1]
 743     return (
 744         isinstance(first, Leaf)
 745         and first.type == token.LPAR
 746         and first.value == ""
 747         and isinstance(last, Leaf)
 748         and last.type == token.RPAR
 749         and last.value == ""
 750     )
 751
 752
 753 def is_empty_par(leaf: Leaf) -> bool:
 754     return is_empty_lpar(leaf) or is_empty_rpar(leaf)
 755
 756
 757 def is_empty_lpar(leaf: Leaf) -> bool:
 758     return leaf.type == token.LPAR and leaf.value == ""
 759
 760
 761 def is_empty_rpar(leaf: Leaf) -> bool:
 762     return leaf.type == token.RPAR and leaf.value == ""
 763
 764
 765 def is_import(leaf: Leaf) -> bool:
 766     """Return True if the given leaf starts an import statement."""
 767     p = leaf.parent
 768     t = leaf.type
 769     v = leaf.value
 770     return bool(
 771         t == token.NAME
 772         and (
 773             (v == "import" and p and p.type == syms.import_name)
 774             or (v == "from" and p and p.type == syms.import_from)
 775         )
 776     )
 777
 778
 779 def is_type_comment(leaf: Leaf, suffix: str = "") -> bool:
 780     """Return True if the given leaf is a special comment.
 781     Only returns true for type comments for now."""
 782     t = leaf.type
 783     v = leaf.value
 784     return t in {token.COMMENT, STANDALONE_COMMENT} and v.startswith("# type:" + suffix)
 785
 786
 787 def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:
 788     """Wrap `child` in parentheses.
 789
 790     This replaces `child` with an atom holding the parentheses and the old
 791     child.  That requires moving the prefix.
 792
 793     If `visible` is False, the leaves will be valueless (and thus invisible).
 794     """
 795     lpar = Leaf(token.LPAR, "(" if visible else "")
 796     rpar = Leaf(token.RPAR, ")" if visible else "")
 797     prefix = child.prefix
 798     child.prefix = ""
 799     index = child.remove() or 0
 800     new_child = Node(syms.atom, [lpar, child, rpar])
 801     new_child.prefix = prefix
 802     parent.insert_child(index, new_child)
 803
 804
 805 def unwrap_singleton_parenthesis(node: LN) -> Optional[LN]:
 806     """Returns `wrapped` if `node` is of the shape ( wrapped ).
 807
 808     Parenthesis can be optional. Returns None otherwise"""
 809     if len(node.children) != 3:
 810         return None
 811
 812     lpar, wrapped, rpar = node.children
 813     if not (lpar.type == token.LPAR and rpar.type == token.RPAR):
 814         return None
 815
 816     return wrapped
 817
 818
 819 def ensure_visible(leaf: Leaf) -> None:
 820     """Make sure parentheses are visible.
 821
 822     They could be invisible as part of some statements (see
 823     :func:`normalize_invisible_parens` and :func:`visit_import_from`).
 824     """
 825     if leaf.type == token.LPAR:
 826         leaf.value = "("
 827     elif leaf.type == token.RPAR:
 828         leaf.value = ")"
 829
 830
 831 def is_name_token(nl: NL) -> TypeGuard[Leaf]:
 832     return nl.type == token.NAME
 833
 834
 835 def is_lpar_token(nl: NL) -> TypeGuard[Leaf]:
 836     return nl.type == token.LPAR
 837
 838
 839 def is_rpar_token(nl: NL) -> TypeGuard[Leaf]:
 840     return nl.type == token.RPAR
 841
 842
 843 def is_string_token(nl: NL) -> TypeGuard[Leaf]:
 844     return nl.type == token.STRING
 845
 846
 847 def is_number_token(nl: NL) -> TypeGuard[Leaf]:
 848     return nl.type == token.NUMBER