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.
17 if sys.version_info >= (3, 8):
18 from typing import Final
20 from typing_extensions import Final
21 if sys.version_info >= (3, 10):
22 from typing import TypeGuard
24 from typing_extensions import TypeGuard
26 from mypy_extensions import mypyc_attr
29 from blib2to3.pytree import Node, Leaf, type_repr, NL
30 from blib2to3 import pygram
31 from blib2to3.pgen2 import token
33 from black.cache import CACHE_DIR
34 from black.strings import has_triple_quotes
37 pygram.initialize(CACHE_DIR)
38 syms: Final = pygram.python_symbols
43 LN = Union[Leaf, Node]
48 WHITESPACE: Final = {token.DEDENT, token.INDENT, token.NEWLINE}
61 STANDALONE_COMMENT: Final = 153
62 token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"
63 LOGIC_OPERATORS: Final = {"and", "or"}
64 COMPARATORS: Final = {
72 MATH_OPERATORS: Final = {
88 STARS: Final = {token.STAR, token.DOUBLESTAR}
89 VARARGS_SPECIALS: Final = STARS | {token.SLASH}
90 VARARGS_PARENTS: Final = {
92 syms.argument, # double star in arglist
93 syms.trailer, # single argument to call
95 syms.varargslist, # lambdas
97 UNPACKING_PARENTS: Final = {
98 syms.atom, # single element of a list or set literal
102 syms.testlist_star_expr,
106 TEST_DESCENDANTS: Final = {
123 ASSIGNMENTS: Final = {
140 IMPLICIT_TUPLE: Final = {syms.testlist, syms.testlist_star_expr, syms.exprlist}
142 token.LPAR: token.RPAR,
143 token.LSQB: token.RSQB,
144 token.LBRACE: token.RBRACE,
146 OPENING_BRACKETS: Final = set(BRACKET.keys())
147 CLOSING_BRACKETS: Final = set(BRACKET.values())
148 BRACKETS: Final = OPENING_BRACKETS | CLOSING_BRACKETS
149 ALWAYS_NO_SPACE: Final = CLOSING_BRACKETS | {token.COMMA, STANDALONE_COMMENT}
154 @mypyc_attr(allow_interpreted_subclasses=True)
155 class Visitor(Generic[T]):
156 """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""
158 def visit(self, node: LN) -> Iterator[T]:
159 """Main method to visit `node` and its children.
161 It tries to find a `visit_*()` method for the given `node.type`, like
162 `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.
163 If no dedicated `visit_*()` method is found, chooses `visit_default()`
166 Then yields objects of type `T` from the selected visitor.
169 name = token.tok_name[node.type]
171 name = str(type_repr(node.type))
172 # We explicitly branch on whether a visitor exists (instead of
173 # using self.visit_default as the default arg to getattr) in order
174 # to save needing to create a bound method object and so mypyc can
175 # generate a native call to visit_default.
176 visitf = getattr(self, f"visit_{name}", None)
178 yield from visitf(node)
180 yield from self.visit_default(node)
182 def visit_default(self, node: LN) -> Iterator[T]:
183 """Default `visit_*()` implementation. Recurses to children of `node`."""
184 if isinstance(node, Node):
185 for child in node.children:
186 yield from self.visit(child)
189 def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901
190 """Return whitespace prefix if needed for the given `leaf`.
192 `complex_subscript` signals whether the given leaf is part of a subscription
193 which has non-trivial arguments, like arithmetic expressions or function calls.
197 DOUBLESPACE: Final = " "
201 if t in ALWAYS_NO_SPACE:
204 if t == token.COMMENT:
207 assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
208 if t == token.COLON and p.type not in {
215 prev = leaf.prev_sibling
217 prevp = preceding_leaf(p)
218 if not prevp or prevp.type in OPENING_BRACKETS:
222 if prevp.type == token.COLON:
225 elif prevp.type != token.COMMA and not complex_subscript:
230 if prevp.type == token.EQUAL:
232 if prevp.parent.type in {
240 elif prevp.parent.type == syms.typedargslist:
241 # A bit hacky: if the equal sign has whitespace, it means we
242 # previously found it's a typed argument. So, we're using
246 elif prevp.type in VARARGS_SPECIALS:
247 if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
250 elif prevp.type == token.COLON:
251 if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:
252 return SPACE if complex_subscript else NO
256 and prevp.parent.type == syms.factor
257 and prevp.type in MATH_OPERATORS
261 elif prevp.type == token.AT and p.parent and p.parent.type == syms.decorator:
262 # no space in decorators
265 elif prev.type in OPENING_BRACKETS:
268 if p.type in {syms.parameters, syms.arglist}:
269 # untyped function signatures or calls
270 if not prev or prev.type != token.COMMA:
273 elif p.type == syms.varargslist:
275 if prev and prev.type != token.COMMA:
278 elif p.type == syms.typedargslist:
279 # typed function signatures
284 if prev.type != syms.tname:
287 elif prev.type == token.EQUAL:
288 # A bit hacky: if the equal sign has whitespace, it means we
289 # previously found it's a typed argument. So, we're using that, too.
292 elif prev.type != token.COMMA:
295 elif p.type == syms.tname:
298 prevp = preceding_leaf(p)
299 if not prevp or prevp.type != token.COMMA:
302 elif p.type == syms.trailer:
303 # attributes and calls
304 if t == token.LPAR or t == token.RPAR:
308 if t == token.DOT or t == token.LSQB:
311 elif prev.type != token.COMMA:
314 elif p.type == syms.argument:
320 prevp = preceding_leaf(p)
321 if not prevp or prevp.type == token.LPAR:
324 elif prev.type in {token.EQUAL} | VARARGS_SPECIALS:
327 elif p.type == syms.decorator:
331 elif p.type == syms.dotted_name:
335 prevp = preceding_leaf(p)
336 if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
339 elif p.type == syms.classdef:
343 if prev and prev.type == token.LPAR:
346 elif p.type in {syms.subscript, syms.sliceop}:
349 assert p.parent is not None, "subscripts are always parented"
350 if p.parent.type == syms.subscriptlist:
355 elif not complex_subscript:
358 elif p.type == syms.atom:
359 if prev and t == token.DOT:
360 # dots, but not the first one.
363 elif p.type == syms.dictsetmaker:
365 if prev and prev.type == token.DOUBLESTAR:
368 elif p.type in {syms.factor, syms.star_expr}:
371 prevp = preceding_leaf(p)
372 if not prevp or prevp.type in OPENING_BRACKETS:
375 prevp_parent = prevp.parent
376 assert prevp_parent is not None
377 if prevp.type == token.COLON and prevp_parent.type in {
383 elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
386 elif t in {token.NAME, token.NUMBER, token.STRING}:
389 elif p.type == syms.import_from:
391 if prev and prev.type == token.DOT:
394 elif t == token.NAME:
398 if prev and prev.type == token.DOT:
401 elif p.type == syms.sliceop:
404 elif p.type == syms.except_clause:
411 def preceding_leaf(node: Optional[LN]) -> Optional[Leaf]:
412 """Return the first leaf that precedes `node`, if any."""
414 res = node.prev_sibling
416 if isinstance(res, Leaf):
420 return list(res.leaves())[-1]
429 def prev_siblings_are(node: Optional[LN], tokens: List[Optional[NodeType]]) -> bool:
430 """Return if the `node` and its previous siblings match types against the provided
431 list of tokens; the provided `node`has its type matched against the last element in
432 the list. `None` can be used as the first element to declare that the start of the
433 list is anchored at the start of its parent's children."""
436 if tokens[-1] is None:
440 if node.type != tokens[-1]:
442 return prev_siblings_are(node.prev_sibling, tokens[:-1])
445 def parent_type(node: Optional[LN]) -> Optional[NodeType]:
448 @node.parent.type, if @node is not None and has a parent.
452 if node is None or node.parent is None:
455 return node.parent.type
458 def child_towards(ancestor: Node, descendant: LN) -> Optional[LN]:
459 """Return the child of `ancestor` that contains `descendant`."""
460 node: Optional[LN] = descendant
461 while node and node.parent != ancestor:
466 def replace_child(old_child: LN, new_child: LN) -> None:
469 * If @old_child.parent is set, replace @old_child with @new_child in
470 @old_child's underlying Node structure.
472 * Otherwise, this function does nothing.
474 parent = old_child.parent
478 child_idx = old_child.remove()
479 if child_idx is not None:
480 parent.insert_child(child_idx, new_child)
483 def container_of(leaf: Leaf) -> LN:
484 """Return `leaf` or one of its ancestors that is the topmost container of it.
486 By "container" we mean a node where `leaf` is the very first child.
488 same_prefix = leaf.prefix
491 parent = container.parent
495 if parent.children[0].prefix != same_prefix:
498 if parent.type == syms.file_input:
501 if parent.prev_sibling is not None and parent.prev_sibling.type in BRACKETS:
508 def first_leaf_column(node: Node) -> Optional[int]:
509 """Returns the column of the first leaf child of a node."""
510 for child in node.children:
511 if isinstance(child, Leaf):
516 def is_arith_like(node: LN) -> bool:
517 """Whether node is an arithmetic or a binary arithmetic expression"""
518 return node.type in {
526 def is_docstring(leaf: Leaf) -> bool:
527 if prev_siblings_are(
528 leaf.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
532 # Multiline docstring on the same line as the `def`.
533 if prev_siblings_are(leaf.parent, [syms.parameters, token.COLON, syms.simple_stmt]):
534 # `syms.parameters` is only used in funcdefs and async_funcdefs in the Python
535 # grammar. We're safe to return True without further checks.
541 def is_empty_tuple(node: LN) -> bool:
542 """Return True if `node` holds an empty tuple."""
544 node.type == syms.atom
545 and len(node.children) == 2
546 and node.children[0].type == token.LPAR
547 and node.children[1].type == token.RPAR
551 def is_one_tuple(node: LN) -> bool:
552 """Return True if `node` holds a tuple with one element, with or without parens."""
553 if node.type == syms.atom:
554 gexp = unwrap_singleton_parenthesis(node)
555 if gexp is None or gexp.type != syms.testlist_gexp:
558 return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA
561 node.type in IMPLICIT_TUPLE
562 and len(node.children) == 2
563 and node.children[1].type == token.COMMA
567 def is_one_sequence_between(
571 brackets: Tuple[int, int] = (token.LPAR, token.RPAR),
573 """Return True if content between `opening` and `closing` is a one-sequence."""
574 if (opening.type, closing.type) != brackets:
577 depth = closing.bracket_depth + 1
578 for _opening_index, leaf in enumerate(leaves):
583 raise LookupError("Opening paren not found in `leaves`")
587 for leaf in leaves[_opening_index:]:
591 bracket_depth = leaf.bracket_depth
592 if bracket_depth == depth and leaf.type == token.COMMA:
594 if leaf.parent and leaf.parent.type in {
604 def is_walrus_assignment(node: LN) -> bool:
605 """Return True iff `node` is of the shape ( test := test )"""
606 inner = unwrap_singleton_parenthesis(node)
607 return inner is not None and inner.type == syms.namedexpr_test
610 def is_simple_decorator_trailer(node: LN, last: bool = False) -> bool:
611 """Return True iff `node` is a trailer valid in a simple decorator"""
612 return node.type == syms.trailer and (
614 len(node.children) == 2
615 and node.children[0].type == token.DOT
616 and node.children[1].type == token.NAME
618 # last trailer can be an argument-less parentheses pair
621 and len(node.children) == 2
622 and node.children[0].type == token.LPAR
623 and node.children[1].type == token.RPAR
625 # last trailer can be arguments
628 and len(node.children) == 3
629 and node.children[0].type == token.LPAR
630 # and node.children[1].type == syms.argument
631 and node.children[2].type == token.RPAR
636 def is_simple_decorator_expression(node: LN) -> bool:
637 """Return True iff `node` could be a 'dotted name' decorator
639 This function takes the node of the 'namedexpr_test' of the new decorator
640 grammar and test if it would be valid under the old decorator grammar.
642 The old grammar was: decorator: @ dotted_name [arguments] NEWLINE
643 The new grammar is : decorator: @ namedexpr_test NEWLINE
645 if node.type == token.NAME:
647 if node.type == syms.power:
650 node.children[0].type == token.NAME
651 and all(map(is_simple_decorator_trailer, node.children[1:-1]))
653 len(node.children) < 2
654 or is_simple_decorator_trailer(node.children[-1], last=True)
660 def is_yield(node: LN) -> bool:
661 """Return True if `node` holds a `yield` or `yield from` expression."""
662 if node.type == syms.yield_expr:
665 if is_name_token(node) and node.value == "yield":
668 if node.type != syms.atom:
671 if len(node.children) != 3:
674 lpar, expr, rpar = node.children
675 if lpar.type == token.LPAR and rpar.type == token.RPAR:
676 return is_yield(expr)
681 def is_vararg(leaf: Leaf, within: Set[NodeType]) -> bool:
682 """Return True if `leaf` is a star or double star in a vararg or kwarg.
684 If `within` includes VARARGS_PARENTS, this applies to function signatures.
685 If `within` includes UNPACKING_PARENTS, it applies to right hand-side
686 extended iterable unpacking (PEP 3132) and additional unpacking
687 generalizations (PEP 448).
689 if leaf.type not in VARARGS_SPECIALS or not leaf.parent:
693 if p.type == syms.star_expr:
694 # Star expressions are also used as assignment targets in extended
695 # iterable unpacking (PEP 3132). See what its parent is instead.
701 return p.type in within
704 def is_multiline_string(leaf: Leaf) -> bool:
705 """Return True if `leaf` is a multiline string that actually spans many lines."""
706 return has_triple_quotes(leaf.value) and "\n" in leaf.value
709 def is_stub_suite(node: Node) -> bool:
710 """Return True if `node` is a suite with a stub body."""
712 len(node.children) != 4
713 or node.children[0].type != token.NEWLINE
714 or node.children[1].type != token.INDENT
715 or node.children[3].type != token.DEDENT
719 return is_stub_body(node.children[2])
722 def is_stub_body(node: LN) -> bool:
723 """Return True if `node` is a simple statement containing an ellipsis."""
724 if not isinstance(node, Node) or node.type != syms.simple_stmt:
727 if len(node.children) != 2:
730 child = node.children[0]
732 child.type == syms.atom
733 and len(child.children) == 3
734 and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
738 def is_atom_with_invisible_parens(node: LN) -> bool:
739 """Given a `LN`, determines whether it's an atom `node` with invisible
740 parens. Useful in dedupe-ing and normalizing parens.
742 if isinstance(node, Leaf) or node.type != syms.atom:
745 first, last = node.children[0], node.children[-1]
747 isinstance(first, Leaf)
748 and first.type == token.LPAR
749 and first.value == ""
750 and isinstance(last, Leaf)
751 and last.type == token.RPAR
756 def is_empty_par(leaf: Leaf) -> bool:
757 return is_empty_lpar(leaf) or is_empty_rpar(leaf)
760 def is_empty_lpar(leaf: Leaf) -> bool:
761 return leaf.type == token.LPAR and leaf.value == ""
764 def is_empty_rpar(leaf: Leaf) -> bool:
765 return leaf.type == token.RPAR and leaf.value == ""
768 def is_import(leaf: Leaf) -> bool:
769 """Return True if the given leaf starts an import statement."""
776 (v == "import" and p and p.type == syms.import_name)
777 or (v == "from" and p and p.type == syms.import_from)
782 def is_type_comment(leaf: Leaf, suffix: str = "") -> bool:
783 """Return True if the given leaf is a special comment.
784 Only returns true for type comments for now."""
787 return t in {token.COMMENT, STANDALONE_COMMENT} and v.startswith("# type:" + suffix)
790 def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:
791 """Wrap `child` in parentheses.
793 This replaces `child` with an atom holding the parentheses and the old
794 child. That requires moving the prefix.
796 If `visible` is False, the leaves will be valueless (and thus invisible).
798 lpar = Leaf(token.LPAR, "(" if visible else "")
799 rpar = Leaf(token.RPAR, ")" if visible else "")
800 prefix = child.prefix
802 index = child.remove() or 0
803 new_child = Node(syms.atom, [lpar, child, rpar])
804 new_child.prefix = prefix
805 parent.insert_child(index, new_child)
808 def unwrap_singleton_parenthesis(node: LN) -> Optional[LN]:
809 """Returns `wrapped` if `node` is of the shape ( wrapped ).
811 Parenthesis can be optional. Returns None otherwise"""
812 if len(node.children) != 3:
815 lpar, wrapped, rpar = node.children
816 if not (lpar.type == token.LPAR and rpar.type == token.RPAR):
822 def ensure_visible(leaf: Leaf) -> None:
823 """Make sure parentheses are visible.
825 They could be invisible as part of some statements (see
826 :func:`normalize_invisible_parens` and :func:`visit_import_from`).
828 if leaf.type == token.LPAR:
830 elif leaf.type == token.RPAR:
834 def is_name_token(nl: NL) -> TypeGuard[Leaf]:
835 return nl.type == token.NAME
838 def is_lpar_token(nl: NL) -> TypeGuard[Leaf]:
839 return nl.type == token.LPAR
842 def is_rpar_token(nl: NL) -> TypeGuard[Leaf]:
843 return nl.type == token.RPAR
846 def is_string_token(nl: NL) -> TypeGuard[Leaf]:
847 return nl.type == token.STRING