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.
18 if sys.version_info < (3, 8):
19 from typing_extensions import Final
21 from typing import Final
24 from blib2to3.pytree import Node, Leaf, type_repr
25 from blib2to3 import pygram
26 from blib2to3.pgen2 import token
28 from black.cache import CACHE_DIR
29 from black.strings import has_triple_quotes
32 pygram.initialize(CACHE_DIR)
33 syms = pygram.python_symbols
38 LN = Union[Leaf, Node]
43 WHITESPACE: Final = {token.DEDENT, token.INDENT, token.NEWLINE}
54 STANDALONE_COMMENT: Final = 153
55 token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"
56 LOGIC_OPERATORS: Final = {"and", "or"}
57 COMPARATORS: Final = {
65 MATH_OPERATORS: Final = {
81 STARS: Final = {token.STAR, token.DOUBLESTAR}
82 VARARGS_SPECIALS: Final = STARS | {token.SLASH}
83 VARARGS_PARENTS: Final = {
85 syms.argument, # double star in arglist
86 syms.trailer, # single argument to call
88 syms.varargslist, # lambdas
90 UNPACKING_PARENTS: Final = {
91 syms.atom, # single element of a list or set literal
95 syms.testlist_star_expr,
97 TEST_DESCENDANTS: Final = {
114 ASSIGNMENTS: Final = {
131 IMPLICIT_TUPLE = {syms.testlist, syms.testlist_star_expr, syms.exprlist}
132 BRACKET = {token.LPAR: token.RPAR, token.LSQB: token.RSQB, token.LBRACE: token.RBRACE}
133 OPENING_BRACKETS = set(BRACKET.keys())
134 CLOSING_BRACKETS = set(BRACKET.values())
135 BRACKETS = OPENING_BRACKETS | CLOSING_BRACKETS
136 ALWAYS_NO_SPACE = CLOSING_BRACKETS | {token.COMMA, STANDALONE_COMMENT}
141 class Visitor(Generic[T]):
142 """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""
144 def visit(self, node: LN) -> Iterator[T]:
145 """Main method to visit `node` and its children.
147 It tries to find a `visit_*()` method for the given `node.type`, like
148 `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.
149 If no dedicated `visit_*()` method is found, chooses `visit_default()`
152 Then yields objects of type `T` from the selected visitor.
155 name = token.tok_name[node.type]
157 name = str(type_repr(node.type))
158 # We explicitly branch on whether a visitor exists (instead of
159 # using self.visit_default as the default arg to getattr) in order
160 # to save needing to create a bound method object and so mypyc can
161 # generate a native call to visit_default.
162 visitf = getattr(self, f"visit_{name}", None)
164 yield from visitf(node)
166 yield from self.visit_default(node)
168 def visit_default(self, node: LN) -> Iterator[T]:
169 """Default `visit_*()` implementation. Recurses to children of `node`."""
170 if isinstance(node, Node):
171 for child in node.children:
172 yield from self.visit(child)
175 def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901
176 """Return whitespace prefix if needed for the given `leaf`.
178 `complex_subscript` signals whether the given leaf is part of a subscription
179 which has non-trivial arguments, like arithmetic expressions or function calls.
187 if t in ALWAYS_NO_SPACE:
190 if t == token.COMMENT:
193 assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
194 if t == token.COLON and p.type not in {
201 prev = leaf.prev_sibling
203 prevp = preceding_leaf(p)
204 if not prevp or prevp.type in OPENING_BRACKETS:
208 if prevp.type == token.COLON:
211 elif prevp.type != token.COMMA and not complex_subscript:
216 if prevp.type == token.EQUAL:
218 if prevp.parent.type in {
226 elif prevp.parent.type == syms.typedargslist:
227 # A bit hacky: if the equal sign has whitespace, it means we
228 # previously found it's a typed argument. So, we're using
232 elif prevp.type in VARARGS_SPECIALS:
233 if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
236 elif prevp.type == token.COLON:
237 if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:
238 return SPACE if complex_subscript else NO
242 and prevp.parent.type == syms.factor
243 and prevp.type in MATH_OPERATORS
248 prevp.type == token.RIGHTSHIFT
250 and prevp.parent.type == syms.shift_expr
251 and prevp.prev_sibling
252 and prevp.prev_sibling.type == token.NAME
253 and prevp.prev_sibling.value == "print" # type: ignore
255 # Python 2 print chevron
257 elif prevp.type == token.AT and p.parent and p.parent.type == syms.decorator:
258 # no space in decorators
261 elif prev.type in OPENING_BRACKETS:
264 if p.type in {syms.parameters, syms.arglist}:
265 # untyped function signatures or calls
266 if not prev or prev.type != token.COMMA:
269 elif p.type == syms.varargslist:
271 if prev and prev.type != token.COMMA:
274 elif p.type == syms.typedargslist:
275 # typed function signatures
280 if prev.type != syms.tname:
283 elif prev.type == token.EQUAL:
284 # A bit hacky: if the equal sign has whitespace, it means we
285 # previously found it's a typed argument. So, we're using that, too.
288 elif prev.type != token.COMMA:
291 elif p.type == syms.tname:
294 prevp = preceding_leaf(p)
295 if not prevp or prevp.type != token.COMMA:
298 elif p.type == syms.trailer:
299 # attributes and calls
300 if t == token.LPAR or t == token.RPAR:
305 prevp = preceding_leaf(p)
306 if not prevp or prevp.type != token.NUMBER:
309 elif t == token.LSQB:
312 elif prev.type != token.COMMA:
315 elif p.type == syms.argument:
321 prevp = preceding_leaf(p)
322 if not prevp or prevp.type == token.LPAR:
325 elif prev.type in {token.EQUAL} | VARARGS_SPECIALS:
328 elif p.type == syms.decorator:
332 elif p.type == syms.dotted_name:
336 prevp = preceding_leaf(p)
337 if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
340 elif p.type == syms.classdef:
344 if prev and prev.type == token.LPAR:
347 elif p.type in {syms.subscript, syms.sliceop}:
350 assert p.parent is not None, "subscripts are always parented"
351 if p.parent.type == syms.subscriptlist:
356 elif not complex_subscript:
359 elif p.type == syms.atom:
360 if prev and t == token.DOT:
361 # dots, but not the first one.
364 elif p.type == syms.dictsetmaker:
366 if prev and prev.type == token.DOUBLESTAR:
369 elif p.type in {syms.factor, syms.star_expr}:
372 prevp = preceding_leaf(p)
373 if not prevp or prevp.type in OPENING_BRACKETS:
376 prevp_parent = prevp.parent
377 assert prevp_parent is not None
378 if prevp.type == token.COLON and prevp_parent.type in {
384 elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
387 elif t in {token.NAME, token.NUMBER, token.STRING}:
390 elif p.type == syms.import_from:
392 if prev and prev.type == token.DOT:
395 elif t == token.NAME:
399 if prev and prev.type == token.DOT:
402 elif p.type == syms.sliceop:
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 last_two_except(leaves: List[Leaf], omit: Collection[LeafID]) -> Tuple[Leaf, Leaf]:
443 """Return (penultimate, last) leaves skipping brackets in `omit` and contents."""
446 for leaf in reversed(leaves):
448 if leaf is stop_after:
456 stop_after = leaf.opening_bracket
460 raise LookupError("Last two leaves were also skipped")
463 def parent_type(node: Optional[LN]) -> Optional[NodeType]:
466 @node.parent.type, if @node is not None and has a parent.
470 if node is None or node.parent is None:
473 return node.parent.type
476 def child_towards(ancestor: Node, descendant: LN) -> Optional[LN]:
477 """Return the child of `ancestor` that contains `descendant`."""
478 node: Optional[LN] = descendant
479 while node and node.parent != ancestor:
484 def replace_child(old_child: LN, new_child: LN) -> None:
487 * If @old_child.parent is set, replace @old_child with @new_child in
488 @old_child's underlying Node structure.
490 * Otherwise, this function does nothing.
492 parent = old_child.parent
496 child_idx = old_child.remove()
497 if child_idx is not None:
498 parent.insert_child(child_idx, new_child)
501 def container_of(leaf: Leaf) -> LN:
502 """Return `leaf` or one of its ancestors that is the topmost container of it.
504 By "container" we mean a node where `leaf` is the very first child.
506 same_prefix = leaf.prefix
509 parent = container.parent
513 if parent.children[0].prefix != same_prefix:
516 if parent.type == syms.file_input:
519 if parent.prev_sibling is not None and parent.prev_sibling.type in BRACKETS:
526 def first_leaf_column(node: Node) -> Optional[int]:
527 """Returns the column of the first leaf child of a node."""
528 for child in node.children:
529 if isinstance(child, Leaf):
534 def first_child_is_arith(node: Node) -> bool:
535 """Whether first child is an arithmetic or a binary arithmetic expression"""
542 return bool(node.children and node.children[0].type in expr_types)
545 def is_docstring(leaf: Leaf) -> bool:
546 if prev_siblings_are(
547 leaf.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
551 # Multiline docstring on the same line as the `def`.
552 if prev_siblings_are(leaf.parent, [syms.parameters, token.COLON, syms.simple_stmt]):
553 # `syms.parameters` is only used in funcdefs and async_funcdefs in the Python
554 # grammar. We're safe to return True without further checks.
560 def is_empty_tuple(node: LN) -> bool:
561 """Return True if `node` holds an empty tuple."""
563 node.type == syms.atom
564 and len(node.children) == 2
565 and node.children[0].type == token.LPAR
566 and node.children[1].type == token.RPAR
570 def is_one_tuple(node: LN) -> bool:
571 """Return True if `node` holds a tuple with one element, with or without parens."""
572 if node.type == syms.atom:
573 gexp = unwrap_singleton_parenthesis(node)
574 if gexp is None or gexp.type != syms.testlist_gexp:
577 return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA
580 node.type in IMPLICIT_TUPLE
581 and len(node.children) == 2
582 and node.children[1].type == token.COMMA
586 def is_one_tuple_between(opening: Leaf, closing: Leaf, leaves: List[Leaf]) -> bool:
587 """Return True if content between `opening` and `closing` looks like a one-tuple."""
588 if opening.type != token.LPAR and closing.type != token.RPAR:
591 depth = closing.bracket_depth + 1
592 for _opening_index, leaf in enumerate(leaves):
597 raise LookupError("Opening paren not found in `leaves`")
601 for leaf in leaves[_opening_index:]:
605 bracket_depth = leaf.bracket_depth
606 if bracket_depth == depth and leaf.type == token.COMMA:
608 if leaf.parent and leaf.parent.type in {
618 def is_walrus_assignment(node: LN) -> bool:
619 """Return True iff `node` is of the shape ( test := test )"""
620 inner = unwrap_singleton_parenthesis(node)
621 return inner is not None and inner.type == syms.namedexpr_test
624 def is_simple_decorator_trailer(node: LN, last: bool = False) -> bool:
625 """Return True iff `node` is a trailer valid in a simple decorator"""
626 return node.type == syms.trailer and (
628 len(node.children) == 2
629 and node.children[0].type == token.DOT
630 and node.children[1].type == token.NAME
632 # last trailer can be an argument-less parentheses pair
635 and len(node.children) == 2
636 and node.children[0].type == token.LPAR
637 and node.children[1].type == token.RPAR
639 # last trailer can be arguments
642 and len(node.children) == 3
643 and node.children[0].type == token.LPAR
644 # and node.children[1].type == syms.argument
645 and node.children[2].type == token.RPAR
650 def is_simple_decorator_expression(node: LN) -> bool:
651 """Return True iff `node` could be a 'dotted name' decorator
653 This function takes the node of the 'namedexpr_test' of the new decorator
654 grammar and test if it would be valid under the old decorator grammar.
656 The old grammar was: decorator: @ dotted_name [arguments] NEWLINE
657 The new grammar is : decorator: @ namedexpr_test NEWLINE
659 if node.type == token.NAME:
661 if node.type == syms.power:
664 node.children[0].type == token.NAME
665 and all(map(is_simple_decorator_trailer, node.children[1:-1]))
667 len(node.children) < 2
668 or is_simple_decorator_trailer(node.children[-1], last=True)
674 def is_yield(node: LN) -> bool:
675 """Return True if `node` holds a `yield` or `yield from` expression."""
676 if node.type == syms.yield_expr:
679 if node.type == token.NAME and node.value == "yield": # type: ignore
682 if node.type != syms.atom:
685 if len(node.children) != 3:
688 lpar, expr, rpar = node.children
689 if lpar.type == token.LPAR and rpar.type == token.RPAR:
690 return is_yield(expr)
695 def is_vararg(leaf: Leaf, within: Set[NodeType]) -> bool:
696 """Return True if `leaf` is a star or double star in a vararg or kwarg.
698 If `within` includes VARARGS_PARENTS, this applies to function signatures.
699 If `within` includes UNPACKING_PARENTS, it applies to right hand-side
700 extended iterable unpacking (PEP 3132) and additional unpacking
701 generalizations (PEP 448).
703 if leaf.type not in VARARGS_SPECIALS or not leaf.parent:
707 if p.type == syms.star_expr:
708 # Star expressions are also used as assignment targets in extended
709 # iterable unpacking (PEP 3132). See what its parent is instead.
715 return p.type in within
718 def is_multiline_string(leaf: Leaf) -> bool:
719 """Return True if `leaf` is a multiline string that actually spans many lines."""
720 return has_triple_quotes(leaf.value) and "\n" in leaf.value
723 def is_stub_suite(node: Node) -> bool:
724 """Return True if `node` is a suite with a stub body."""
726 len(node.children) != 4
727 or node.children[0].type != token.NEWLINE
728 or node.children[1].type != token.INDENT
729 or node.children[3].type != token.DEDENT
733 return is_stub_body(node.children[2])
736 def is_stub_body(node: LN) -> bool:
737 """Return True if `node` is a simple statement containing an ellipsis."""
738 if not isinstance(node, Node) or node.type != syms.simple_stmt:
741 if len(node.children) != 2:
744 child = node.children[0]
746 child.type == syms.atom
747 and len(child.children) == 3
748 and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
752 def is_atom_with_invisible_parens(node: LN) -> bool:
753 """Given a `LN`, determines whether it's an atom `node` with invisible
754 parens. Useful in dedupe-ing and normalizing parens.
756 if isinstance(node, Leaf) or node.type != syms.atom:
759 first, last = node.children[0], node.children[-1]
761 isinstance(first, Leaf)
762 and first.type == token.LPAR
763 and first.value == ""
764 and isinstance(last, Leaf)
765 and last.type == token.RPAR
770 def is_empty_par(leaf: Leaf) -> bool:
771 return is_empty_lpar(leaf) or is_empty_rpar(leaf)
774 def is_empty_lpar(leaf: Leaf) -> bool:
775 return leaf.type == token.LPAR and leaf.value == ""
778 def is_empty_rpar(leaf: Leaf) -> bool:
779 return leaf.type == token.RPAR and leaf.value == ""
782 def is_import(leaf: Leaf) -> bool:
783 """Return True if the given leaf starts an import statement."""
790 (v == "import" and p and p.type == syms.import_name)
791 or (v == "from" and p and p.type == syms.import_from)
796 def is_type_comment(leaf: Leaf, suffix: str = "") -> bool:
797 """Return True if the given leaf is a special comment.
798 Only returns true for type comments for now."""
801 return t in {token.COMMENT, STANDALONE_COMMENT} and v.startswith("# type:" + suffix)
804 def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:
805 """Wrap `child` in parentheses.
807 This replaces `child` with an atom holding the parentheses and the old
808 child. That requires moving the prefix.
810 If `visible` is False, the leaves will be valueless (and thus invisible).
812 lpar = Leaf(token.LPAR, "(" if visible else "")
813 rpar = Leaf(token.RPAR, ")" if visible else "")
814 prefix = child.prefix
816 index = child.remove() or 0
817 new_child = Node(syms.atom, [lpar, child, rpar])
818 new_child.prefix = prefix
819 parent.insert_child(index, new_child)
822 def unwrap_singleton_parenthesis(node: LN) -> Optional[LN]:
823 """Returns `wrapped` if `node` is of the shape ( wrapped ).
825 Parenthesis can be optional. Returns None otherwise"""
826 if len(node.children) != 3:
829 lpar, wrapped, rpar = node.children
830 if not (lpar.type == token.LPAR and rpar.type == token.RPAR):
836 def ensure_visible(leaf: Leaf) -> None:
837 """Make sure parentheses are visible.
839 They could be invisible as part of some statements (see
840 :func:`normalize_invisible_parens` and :func:`visit_import_from`).
842 if leaf.type == token.LPAR:
844 elif leaf.type == token.RPAR: