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.
   2 blib2to3 Node/Leaf transformation-related utility functions.
 
   6 from typing import Generic, Iterator, List, Optional, Set, Tuple, TypeVar, Union
 
   8 if sys.version_info >= (3, 8):
 
   9     from typing import Final
 
  11     from typing_extensions import Final
 
  12 if sys.version_info >= (3, 10):
 
  13     from typing import TypeGuard
 
  15     from typing_extensions import TypeGuard
 
  17 from mypy_extensions import mypyc_attr
 
  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
 
  25 pygram.initialize(CACHE_DIR)
 
  26 syms: Final = pygram.python_symbols
 
  31 LN = Union[Leaf, Node]
 
  36 WHITESPACE: Final = {token.DEDENT, token.INDENT, token.NEWLINE}
 
  49 STANDALONE_COMMENT: Final = 153
 
  50 token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"
 
  51 LOGIC_OPERATORS: Final = {"and", "or"}
 
  52 COMPARATORS: Final = {
 
  60 MATH_OPERATORS: Final = {
 
  76 STARS: Final = {token.STAR, token.DOUBLESTAR}
 
  77 VARARGS_SPECIALS: Final = STARS | {token.SLASH}
 
  78 VARARGS_PARENTS: Final = {
 
  80     syms.argument,  # double star in arglist
 
  81     syms.trailer,  # single argument to call
 
  83     syms.varargslist,  # lambdas
 
  85 UNPACKING_PARENTS: Final = {
 
  86     syms.atom,  # single element of a list or set literal
 
  90     syms.testlist_star_expr,
 
  94 TEST_DESCENDANTS: Final = {
 
 111 TYPED_NAMES: Final = {syms.tname, syms.tname_star}
 
 112 ASSIGNMENTS: Final = {
 
 129 IMPLICIT_TUPLE: Final = {syms.testlist, syms.testlist_star_expr, syms.exprlist}
 
 131     token.LPAR: token.RPAR,
 
 132     token.LSQB: token.RSQB,
 
 133     token.LBRACE: token.RBRACE,
 
 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}
 
 143 @mypyc_attr(allow_interpreted_subclasses=True)
 
 144 class Visitor(Generic[T]):
 
 145     """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""
 
 147     def visit(self, node: LN) -> Iterator[T]:
 
 148         """Main method to visit `node` and its children.
 
 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()`
 
 155         Then yields objects of type `T` from the selected visitor.
 
 158             name = token.tok_name[node.type]
 
 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)
 
 167             yield from visitf(node)
 
 169             yield from self.visit_default(node)
 
 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)
 
 178 def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str:  # noqa: C901
 
 179     """Return whitespace prefix if needed for the given `leaf`.
 
 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.
 
 186     DOUBLESPACE: Final = "  "
 
 190     if t in ALWAYS_NO_SPACE:
 
 193     if t == token.COMMENT:
 
 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 {
 
 204     prev = leaf.prev_sibling
 
 206         prevp = preceding_leaf(p)
 
 207         if not prevp or prevp.type in OPENING_BRACKETS:
 
 211             if prevp.type == token.COLON:
 
 214             elif prevp.type != token.COMMA and not complex_subscript:
 
 219         if prevp.type == token.EQUAL:
 
 221                 if prevp.parent.type in {
 
 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
 
 236             prevp.type == token.STAR
 
 237             and parent_type(prevp) == syms.star_expr
 
 238             and parent_type(prevp.parent) == syms.subscriptlist
 
 240             # No space between typevar tuples.
 
 243         elif prevp.type in VARARGS_SPECIALS:
 
 244             if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
 
 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
 
 253             and prevp.parent.type == syms.factor
 
 254             and prevp.type in MATH_OPERATORS
 
 258         elif prevp.type == token.AT and p.parent and p.parent.type == syms.decorator:
 
 259             # no space in decorators
 
 262     elif prev.type in OPENING_BRACKETS:
 
 265     if p.type in {syms.parameters, syms.arglist}:
 
 266         # untyped function signatures or calls
 
 267         if not prev or prev.type != token.COMMA:
 
 270     elif p.type == syms.varargslist:
 
 272         if prev and prev.type != token.COMMA:
 
 275     elif p.type == syms.typedargslist:
 
 276         # typed function signatures
 
 281             if prev.type not in TYPED_NAMES:
 
 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.
 
 289         elif prev.type != token.COMMA:
 
 292     elif p.type in TYPED_NAMES:
 
 295             prevp = preceding_leaf(p)
 
 296             if not prevp or prevp.type != token.COMMA:
 
 299     elif p.type == syms.trailer:
 
 300         # attributes and calls
 
 301         if t == token.LPAR or t == token.RPAR:
 
 305             if t == token.DOT or t == token.LSQB:
 
 308         elif prev.type != token.COMMA:
 
 311     elif p.type == syms.argument:
 
 317             prevp = preceding_leaf(p)
 
 318             if not prevp or prevp.type == token.LPAR:
 
 321         elif prev.type in {token.EQUAL} | VARARGS_SPECIALS:
 
 324     elif p.type == syms.decorator:
 
 328     elif p.type == syms.dotted_name:
 
 332         prevp = preceding_leaf(p)
 
 333         if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
 
 336     elif p.type == syms.classdef:
 
 340         if prev and prev.type == token.LPAR:
 
 343     elif p.type in {syms.subscript, syms.sliceop}:
 
 346             assert p.parent is not None, "subscripts are always parented"
 
 347             if p.parent.type == syms.subscriptlist:
 
 352         elif not complex_subscript:
 
 355     elif p.type == syms.atom:
 
 356         if prev and t == token.DOT:
 
 357             # dots, but not the first one.
 
 360     elif p.type == syms.dictsetmaker:
 
 362         if prev and prev.type == token.DOUBLESTAR:
 
 365     elif p.type in {syms.factor, syms.star_expr}:
 
 368             prevp = preceding_leaf(p)
 
 369             if not prevp or prevp.type in OPENING_BRACKETS:
 
 372             prevp_parent = prevp.parent
 
 373             assert prevp_parent is not None
 
 374             if prevp.type == token.COLON and prevp_parent.type in {
 
 380             elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
 
 383         elif t in {token.NAME, token.NUMBER, token.STRING}:
 
 386     elif p.type == syms.import_from:
 
 388             if prev and prev.type == token.DOT:
 
 391         elif t == token.NAME:
 
 395             if prev and prev.type == token.DOT:
 
 398     elif p.type == syms.sliceop:
 
 401     elif p.type == syms.except_clause:
 
 408 def preceding_leaf(node: Optional[LN]) -> Optional[Leaf]:
 
 409     """Return the first leaf that precedes `node`, if any."""
 
 411         res = node.prev_sibling
 
 413             if isinstance(res, Leaf):
 
 417                 return list(res.leaves())[-1]
 
 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."""
 
 433     if tokens[-1] is None:
 
 437     if node.type != tokens[-1]:
 
 439     return prev_siblings_are(node.prev_sibling, tokens[:-1])
 
 442 def parent_type(node: Optional[LN]) -> Optional[NodeType]:
 
 445         @node.parent.type, if @node is not None and has a parent.
 
 449     if node is None or node.parent is None:
 
 452     return node.parent.type
 
 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:
 
 463 def replace_child(old_child: LN, new_child: LN) -> None:
 
 466         * If @old_child.parent is set, replace @old_child with @new_child in
 
 467         @old_child's underlying Node structure.
 
 469         * Otherwise, this function does nothing.
 
 471     parent = old_child.parent
 
 475     child_idx = old_child.remove()
 
 476     if child_idx is not None:
 
 477         parent.insert_child(child_idx, new_child)
 
 480 def container_of(leaf: Leaf) -> LN:
 
 481     """Return `leaf` or one of its ancestors that is the topmost container of it.
 
 483     By "container" we mean a node where `leaf` is the very first child.
 
 485     same_prefix = leaf.prefix
 
 488         parent = container.parent
 
 492         if parent.children[0].prefix != same_prefix:
 
 495         if parent.type == syms.file_input:
 
 498         if parent.prev_sibling is not None and parent.prev_sibling.type in BRACKETS:
 
 505 def first_leaf_of(node: LN) -> Optional[Leaf]:
 
 506     """Returns the first leaf of the node tree."""
 
 507     if isinstance(node, Leaf):
 
 510         return first_leaf_of(node.children[0])
 
 515 def is_arith_like(node: LN) -> bool:
 
 516     """Whether node is an arithmetic or a binary arithmetic expression"""
 
 517     return node.type in {
 
 525 def is_docstring(leaf: Leaf) -> bool:
 
 526     if prev_siblings_are(
 
 527         leaf.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
 
 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.
 
 540 def is_empty_tuple(node: LN) -> bool:
 
 541     """Return True if `node` holds an empty tuple."""
 
 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
 
 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:
 
 557         return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA
 
 560         node.type in IMPLICIT_TUPLE
 
 561         and len(node.children) == 2
 
 562         and node.children[1].type == token.COMMA
 
 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:
 
 570     gexp = unwrap_singleton_parenthesis(node)
 
 571     if gexp is None or gexp.type != syms.testlist_gexp:
 
 574     return any(child.type == syms.namedexpr_test for child in gexp.children)
 
 577 def is_one_sequence_between(
 
 581     brackets: Tuple[int, int] = (token.LPAR, token.RPAR),
 
 583     """Return True if content between `opening` and `closing` is a one-sequence."""
 
 584     if (opening.type, closing.type) != brackets:
 
 587     depth = closing.bracket_depth + 1
 
 588     for _opening_index, leaf in enumerate(leaves):
 
 593         raise LookupError("Opening paren not found in `leaves`")
 
 597     for leaf in leaves[_opening_index:]:
 
 601         bracket_depth = leaf.bracket_depth
 
 602         if bracket_depth == depth and leaf.type == token.COMMA:
 
 604             if leaf.parent and leaf.parent.type in {
 
 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
 
 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 (
 
 624             len(node.children) == 2
 
 625             and node.children[0].type == token.DOT
 
 626             and node.children[1].type == token.NAME
 
 628         # last trailer can be an argument-less parentheses pair
 
 631             and len(node.children) == 2
 
 632             and node.children[0].type == token.LPAR
 
 633             and node.children[1].type == token.RPAR
 
 635         # last trailer can be arguments
 
 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
 
 646 def is_simple_decorator_expression(node: LN) -> bool:
 
 647     """Return True iff `node` could be a 'dotted name' decorator
 
 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.
 
 652     The old grammar was: decorator: @ dotted_name [arguments] NEWLINE
 
 653     The new grammar is : decorator: @ namedexpr_test NEWLINE
 
 655     if node.type == token.NAME:
 
 657     if node.type == syms.power:
 
 660                 node.children[0].type == token.NAME
 
 661                 and all(map(is_simple_decorator_trailer, node.children[1:-1]))
 
 663                     len(node.children) < 2
 
 664                     or is_simple_decorator_trailer(node.children[-1], last=True)
 
 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:
 
 675     if is_name_token(node) and node.value == "yield":
 
 678     if node.type != syms.atom:
 
 681     if len(node.children) != 3:
 
 684     lpar, expr, rpar = node.children
 
 685     if lpar.type == token.LPAR and rpar.type == token.RPAR:
 
 686         return is_yield(expr)
 
 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.
 
 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).
 
 699     if leaf.type not in VARARGS_SPECIALS or not 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.
 
 711     return p.type in within
 
 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
 
 719 def is_stub_suite(node: Node) -> bool:
 
 720     """Return True if `node` is a suite with a stub body."""
 
 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
 
 729     return is_stub_body(node.children[2])
 
 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:
 
 737     if len(node.children) != 2:
 
 740     child = node.children[0]
 
 742         child.type == syms.atom
 
 743         and len(child.children) == 3
 
 744         and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
 
 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.
 
 752     if isinstance(node, Leaf) or node.type != syms.atom:
 
 755     first, last = node.children[0], node.children[-1]
 
 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
 
 766 def is_empty_par(leaf: Leaf) -> bool:
 
 767     return is_empty_lpar(leaf) or is_empty_rpar(leaf)
 
 770 def is_empty_lpar(leaf: Leaf) -> bool:
 
 771     return leaf.type == token.LPAR and leaf.value == ""
 
 774 def is_empty_rpar(leaf: Leaf) -> bool:
 
 775     return leaf.type == token.RPAR and leaf.value == ""
 
 778 def is_import(leaf: Leaf) -> bool:
 
 779     """Return True if the given leaf starts an import statement."""
 
 786             (v == "import" and p and p.type == syms.import_name)
 
 787             or (v == "from" and p and p.type == syms.import_from)
 
 792 def is_type_comment(leaf: Leaf, suffix: str = "") -> bool:
 
 793     """Return True if the given leaf is a special comment.
 
 794     Only returns true for type comments for now."""
 
 797     return t in {token.COMMENT, STANDALONE_COMMENT} and v.startswith("# type:" + suffix)
 
 800 def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:
 
 801     """Wrap `child` in parentheses.
 
 803     This replaces `child` with an atom holding the parentheses and the old
 
 804     child.  That requires moving the prefix.
 
 806     If `visible` is False, the leaves will be valueless (and thus invisible).
 
 808     lpar = Leaf(token.LPAR, "(" if visible else "")
 
 809     rpar = Leaf(token.RPAR, ")" if visible else "")
 
 810     prefix = child.prefix
 
 812     index = child.remove() or 0
 
 813     new_child = Node(syms.atom, [lpar, child, rpar])
 
 814     new_child.prefix = prefix
 
 815     parent.insert_child(index, new_child)
 
 818 def unwrap_singleton_parenthesis(node: LN) -> Optional[LN]:
 
 819     """Returns `wrapped` if `node` is of the shape ( wrapped ).
 
 821     Parenthesis can be optional. Returns None otherwise"""
 
 822     if len(node.children) != 3:
 
 825     lpar, wrapped, rpar = node.children
 
 826     if not (lpar.type == token.LPAR and rpar.type == token.RPAR):
 
 832 def ensure_visible(leaf: Leaf) -> None:
 
 833     """Make sure parentheses are visible.
 
 835     They could be invisible as part of some statements (see
 
 836     :func:`normalize_invisible_parens` and :func:`visit_import_from`).
 
 838     if leaf.type == token.LPAR:
 
 840     elif leaf.type == token.RPAR:
 
 844 def is_name_token(nl: NL) -> TypeGuard[Leaf]:
 
 845     return nl.type == token.NAME
 
 848 def is_lpar_token(nl: NL) -> TypeGuard[Leaf]:
 
 849     return nl.type == token.LPAR
 
 852 def is_rpar_token(nl: NL) -> TypeGuard[Leaf]:
 
 853     return nl.type == token.RPAR
 
 856 def is_string_token(nl: NL) -> TypeGuard[Leaf]:
 
 857     return nl.type == token.STRING
 
 860 def is_number_token(nl: NL) -> TypeGuard[Leaf]:
 
 861     return nl.type == token.NUMBER
 
 864 def is_part_of_annotation(leaf: Leaf) -> bool:
 
 865     """Returns whether this leaf is part of type annotations."""
 
 866     ancestor = leaf.parent
 
 867     while ancestor is not None:
 
 868         if ancestor.prev_sibling and ancestor.prev_sibling.type == token.RARROW:
 
 870         if ancestor.parent and ancestor.parent.type == syms.tname:
 
 872         ancestor = ancestor.parent