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 Final, Generic, Iterator, List, Optional, Set, Tuple, TypeVar, Union
8 if sys.version_info >= (3, 10):
9 from typing import TypeGuard
11 from typing_extensions import TypeGuard
13 from mypy_extensions import mypyc_attr
15 from black.cache import CACHE_DIR
16 from black.strings import has_triple_quotes
17 from blib2to3 import pygram
18 from blib2to3.pgen2 import token
19 from blib2to3.pytree import NL, Leaf, Node, type_repr
21 pygram.initialize(CACHE_DIR)
22 syms: Final = pygram.python_symbols
27 LN = Union[Leaf, Node]
32 WHITESPACE: Final = {token.DEDENT, token.INDENT, token.NEWLINE}
45 STANDALONE_COMMENT: Final = 153
46 token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"
47 LOGIC_OPERATORS: Final = {"and", "or"}
48 COMPARATORS: Final = {
56 MATH_OPERATORS: Final = {
72 STARS: Final = {token.STAR, token.DOUBLESTAR}
73 VARARGS_SPECIALS: Final = STARS | {token.SLASH}
74 VARARGS_PARENTS: Final = {
76 syms.argument, # double star in arglist
77 syms.trailer, # single argument to call
79 syms.varargslist, # lambdas
81 UNPACKING_PARENTS: Final = {
82 syms.atom, # single element of a list or set literal
86 syms.testlist_star_expr,
90 TEST_DESCENDANTS: Final = {
107 TYPED_NAMES: Final = {syms.tname, syms.tname_star}
108 ASSIGNMENTS: Final = {
125 IMPLICIT_TUPLE: Final = {syms.testlist, syms.testlist_star_expr, syms.exprlist}
127 token.LPAR: token.RPAR,
128 token.LSQB: token.RSQB,
129 token.LBRACE: token.RBRACE,
131 OPENING_BRACKETS: Final = set(BRACKET.keys())
132 CLOSING_BRACKETS: Final = set(BRACKET.values())
133 BRACKETS: Final = OPENING_BRACKETS | CLOSING_BRACKETS
134 ALWAYS_NO_SPACE: Final = CLOSING_BRACKETS | {token.COMMA, STANDALONE_COMMENT}
139 @mypyc_attr(allow_interpreted_subclasses=True)
140 class Visitor(Generic[T]):
141 """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""
143 def visit(self, node: LN) -> Iterator[T]:
144 """Main method to visit `node` and its children.
146 It tries to find a `visit_*()` method for the given `node.type`, like
147 `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.
148 If no dedicated `visit_*()` method is found, chooses `visit_default()`
151 Then yields objects of type `T` from the selected visitor.
154 name = token.tok_name[node.type]
156 name = str(type_repr(node.type))
157 # We explicitly branch on whether a visitor exists (instead of
158 # using self.visit_default as the default arg to getattr) in order
159 # to save needing to create a bound method object and so mypyc can
160 # generate a native call to visit_default.
161 visitf = getattr(self, f"visit_{name}", None)
163 yield from visitf(node)
165 yield from self.visit_default(node)
167 def visit_default(self, node: LN) -> Iterator[T]:
168 """Default `visit_*()` implementation. Recurses to children of `node`."""
169 if isinstance(node, Node):
170 for child in node.children:
171 yield from self.visit(child)
174 def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901
175 """Return whitespace prefix if needed for the given `leaf`.
177 `complex_subscript` signals whether the given leaf is part of a subscription
178 which has non-trivial arguments, like arithmetic expressions or function calls.
181 SPACE: Final[str] = " "
182 DOUBLESPACE: Final[str] = " "
186 if t in ALWAYS_NO_SPACE:
189 if t == token.COMMENT:
192 assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
193 if t == token.COLON and p.type not in {
200 prev = leaf.prev_sibling
202 prevp = preceding_leaf(p)
203 if not prevp or prevp.type in OPENING_BRACKETS:
207 if prevp.type == token.COLON:
210 elif prevp.type != token.COMMA and not complex_subscript:
215 if prevp.type == token.EQUAL:
217 if prevp.parent.type in {
225 elif prevp.parent.type == syms.typedargslist:
226 # A bit hacky: if the equal sign has whitespace, it means we
227 # previously found it's a typed argument. So, we're using
232 prevp.type == token.STAR
233 and parent_type(prevp) == syms.star_expr
234 and parent_type(prevp.parent) == syms.subscriptlist
236 # No space between typevar tuples.
239 elif prevp.type in VARARGS_SPECIALS:
240 if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
243 elif prevp.type == token.COLON:
244 if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:
245 return SPACE if complex_subscript else NO
249 and prevp.parent.type == syms.factor
250 and prevp.type in MATH_OPERATORS
254 elif prevp.type == token.AT and p.parent and p.parent.type == syms.decorator:
255 # no space in decorators
258 elif prev.type in OPENING_BRACKETS:
261 if p.type in {syms.parameters, syms.arglist}:
262 # untyped function signatures or calls
263 if not prev or prev.type != token.COMMA:
266 elif p.type == syms.varargslist:
268 if prev and prev.type != token.COMMA:
271 elif p.type == syms.typedargslist:
272 # typed function signatures
277 if prev.type not in TYPED_NAMES:
280 elif prev.type == token.EQUAL:
281 # A bit hacky: if the equal sign has whitespace, it means we
282 # previously found it's a typed argument. So, we're using that, too.
285 elif prev.type != token.COMMA:
288 elif p.type in TYPED_NAMES:
291 prevp = preceding_leaf(p)
292 if not prevp or prevp.type != token.COMMA:
295 elif p.type == syms.trailer:
296 # attributes and calls
297 if t == token.LPAR or t == token.RPAR:
301 if t == token.DOT or t == token.LSQB:
304 elif prev.type != token.COMMA:
307 elif p.type == syms.argument:
313 prevp = preceding_leaf(p)
314 if not prevp or prevp.type == token.LPAR:
317 elif prev.type in {token.EQUAL} | VARARGS_SPECIALS:
320 elif p.type == syms.decorator:
324 elif p.type == syms.dotted_name:
328 prevp = preceding_leaf(p)
329 if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
332 elif p.type == syms.classdef:
336 if prev and prev.type == token.LPAR:
339 elif p.type in {syms.subscript, syms.sliceop}:
342 assert p.parent is not None, "subscripts are always parented"
343 if p.parent.type == syms.subscriptlist:
348 elif not complex_subscript:
351 elif p.type == syms.atom:
352 if prev and t == token.DOT:
353 # dots, but not the first one.
356 elif p.type == syms.dictsetmaker:
358 if prev and prev.type == token.DOUBLESTAR:
361 elif p.type in {syms.factor, syms.star_expr}:
364 prevp = preceding_leaf(p)
365 if not prevp or prevp.type in OPENING_BRACKETS:
368 prevp_parent = prevp.parent
369 assert prevp_parent is not None
370 if prevp.type == token.COLON and prevp_parent.type in {
376 elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
379 elif t in {token.NAME, token.NUMBER, token.STRING}:
382 elif p.type == syms.import_from:
384 if prev and prev.type == token.DOT:
387 elif t == token.NAME:
391 if prev and prev.type == token.DOT:
394 elif p.type == syms.sliceop:
397 elif p.type == syms.except_clause:
404 def preceding_leaf(node: Optional[LN]) -> Optional[Leaf]:
405 """Return the first leaf that precedes `node`, if any."""
407 res = node.prev_sibling
409 if isinstance(res, Leaf):
413 return list(res.leaves())[-1]
422 def prev_siblings_are(node: Optional[LN], tokens: List[Optional[NodeType]]) -> bool:
423 """Return if the `node` and its previous siblings match types against the provided
424 list of tokens; the provided `node`has its type matched against the last element in
425 the list. `None` can be used as the first element to declare that the start of the
426 list is anchored at the start of its parent's children."""
429 if tokens[-1] is None:
433 if node.type != tokens[-1]:
435 return prev_siblings_are(node.prev_sibling, tokens[:-1])
438 def parent_type(node: Optional[LN]) -> Optional[NodeType]:
441 @node.parent.type, if @node is not None and has a parent.
445 if node is None or node.parent is None:
448 return node.parent.type
451 def child_towards(ancestor: Node, descendant: LN) -> Optional[LN]:
452 """Return the child of `ancestor` that contains `descendant`."""
453 node: Optional[LN] = descendant
454 while node and node.parent != ancestor:
459 def replace_child(old_child: LN, new_child: LN) -> None:
462 * If @old_child.parent is set, replace @old_child with @new_child in
463 @old_child's underlying Node structure.
465 * Otherwise, this function does nothing.
467 parent = old_child.parent
471 child_idx = old_child.remove()
472 if child_idx is not None:
473 parent.insert_child(child_idx, new_child)
476 def container_of(leaf: Leaf) -> LN:
477 """Return `leaf` or one of its ancestors that is the topmost container of it.
479 By "container" we mean a node where `leaf` is the very first child.
481 same_prefix = leaf.prefix
484 parent = container.parent
488 if parent.children[0].prefix != same_prefix:
491 if parent.type == syms.file_input:
494 if parent.prev_sibling is not None and parent.prev_sibling.type in BRACKETS:
501 def first_leaf_of(node: LN) -> Optional[Leaf]:
502 """Returns the first leaf of the node tree."""
503 if isinstance(node, Leaf):
506 return first_leaf_of(node.children[0])
511 def is_arith_like(node: LN) -> bool:
512 """Whether node is an arithmetic or a binary arithmetic expression"""
513 return node.type in {
521 def is_docstring(leaf: Leaf) -> bool:
522 if prev_siblings_are(
523 leaf.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
527 # Multiline docstring on the same line as the `def`.
528 if prev_siblings_are(leaf.parent, [syms.parameters, token.COLON, syms.simple_stmt]):
529 # `syms.parameters` is only used in funcdefs and async_funcdefs in the Python
530 # grammar. We're safe to return True without further checks.
536 def is_empty_tuple(node: LN) -> bool:
537 """Return True if `node` holds an empty tuple."""
539 node.type == syms.atom
540 and len(node.children) == 2
541 and node.children[0].type == token.LPAR
542 and node.children[1].type == token.RPAR
546 def is_one_tuple(node: LN) -> bool:
547 """Return True if `node` holds a tuple with one element, with or without parens."""
548 if node.type == syms.atom:
549 gexp = unwrap_singleton_parenthesis(node)
550 if gexp is None or gexp.type != syms.testlist_gexp:
553 return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA
556 node.type in IMPLICIT_TUPLE
557 and len(node.children) == 2
558 and node.children[1].type == token.COMMA
562 def is_tuple_containing_walrus(node: LN) -> bool:
563 """Return True if `node` holds a tuple that contains a walrus operator."""
564 if node.type != syms.atom:
566 gexp = unwrap_singleton_parenthesis(node)
567 if gexp is None or gexp.type != syms.testlist_gexp:
570 return any(child.type == syms.namedexpr_test for child in gexp.children)
573 def is_one_sequence_between(
577 brackets: Tuple[int, int] = (token.LPAR, token.RPAR),
579 """Return True if content between `opening` and `closing` is a one-sequence."""
580 if (opening.type, closing.type) != brackets:
583 depth = closing.bracket_depth + 1
584 for _opening_index, leaf in enumerate(leaves):
589 raise LookupError("Opening paren not found in `leaves`")
593 for leaf in leaves[_opening_index:]:
597 bracket_depth = leaf.bracket_depth
598 if bracket_depth == depth and leaf.type == token.COMMA:
600 if leaf.parent and leaf.parent.type in {
610 def is_walrus_assignment(node: LN) -> bool:
611 """Return True iff `node` is of the shape ( test := test )"""
612 inner = unwrap_singleton_parenthesis(node)
613 return inner is not None and inner.type == syms.namedexpr_test
616 def is_simple_decorator_trailer(node: LN, last: bool = False) -> bool:
617 """Return True iff `node` is a trailer valid in a simple decorator"""
618 return node.type == syms.trailer and (
620 len(node.children) == 2
621 and node.children[0].type == token.DOT
622 and node.children[1].type == token.NAME
624 # last trailer can be an argument-less parentheses pair
627 and len(node.children) == 2
628 and node.children[0].type == token.LPAR
629 and node.children[1].type == token.RPAR
631 # last trailer can be arguments
634 and len(node.children) == 3
635 and node.children[0].type == token.LPAR
636 # and node.children[1].type == syms.argument
637 and node.children[2].type == token.RPAR
642 def is_simple_decorator_expression(node: LN) -> bool:
643 """Return True iff `node` could be a 'dotted name' decorator
645 This function takes the node of the 'namedexpr_test' of the new decorator
646 grammar and test if it would be valid under the old decorator grammar.
648 The old grammar was: decorator: @ dotted_name [arguments] NEWLINE
649 The new grammar is : decorator: @ namedexpr_test NEWLINE
651 if node.type == token.NAME:
653 if node.type == syms.power:
656 node.children[0].type == token.NAME
657 and all(map(is_simple_decorator_trailer, node.children[1:-1]))
659 len(node.children) < 2
660 or is_simple_decorator_trailer(node.children[-1], last=True)
666 def is_yield(node: LN) -> bool:
667 """Return True if `node` holds a `yield` or `yield from` expression."""
668 if node.type == syms.yield_expr:
671 if is_name_token(node) and node.value == "yield":
674 if node.type != syms.atom:
677 if len(node.children) != 3:
680 lpar, expr, rpar = node.children
681 if lpar.type == token.LPAR and rpar.type == token.RPAR:
682 return is_yield(expr)
687 def is_vararg(leaf: Leaf, within: Set[NodeType]) -> bool:
688 """Return True if `leaf` is a star or double star in a vararg or kwarg.
690 If `within` includes VARARGS_PARENTS, this applies to function signatures.
691 If `within` includes UNPACKING_PARENTS, it applies to right hand-side
692 extended iterable unpacking (PEP 3132) and additional unpacking
693 generalizations (PEP 448).
695 if leaf.type not in VARARGS_SPECIALS or not leaf.parent:
699 if p.type == syms.star_expr:
700 # Star expressions are also used as assignment targets in extended
701 # iterable unpacking (PEP 3132). See what its parent is instead.
707 return p.type in within
710 def is_multiline_string(leaf: Leaf) -> bool:
711 """Return True if `leaf` is a multiline string that actually spans many lines."""
712 return has_triple_quotes(leaf.value) and "\n" in leaf.value
715 def is_stub_suite(node: Node) -> bool:
716 """Return True if `node` is a suite with a stub body."""
718 # If there is a comment, we want to keep it.
719 if node.prefix.strip():
723 len(node.children) != 4
724 or node.children[0].type != token.NEWLINE
725 or node.children[1].type != token.INDENT
726 or node.children[3].type != token.DEDENT
730 if node.children[3].prefix.strip():
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 not child.prefix.strip()
747 and child.type == syms.atom
748 and len(child.children) == 3
749 and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
753 def is_atom_with_invisible_parens(node: LN) -> bool:
754 """Given a `LN`, determines whether it's an atom `node` with invisible
755 parens. Useful in dedupe-ing and normalizing parens.
757 if isinstance(node, Leaf) or node.type != syms.atom:
760 first, last = node.children[0], node.children[-1]
762 isinstance(first, Leaf)
763 and first.type == token.LPAR
764 and first.value == ""
765 and isinstance(last, Leaf)
766 and last.type == token.RPAR
771 def is_empty_par(leaf: Leaf) -> bool:
772 return is_empty_lpar(leaf) or is_empty_rpar(leaf)
775 def is_empty_lpar(leaf: Leaf) -> bool:
776 return leaf.type == token.LPAR and leaf.value == ""
779 def is_empty_rpar(leaf: Leaf) -> bool:
780 return leaf.type == token.RPAR and leaf.value == ""
783 def is_import(leaf: Leaf) -> bool:
784 """Return True if the given leaf starts an import statement."""
791 (v == "import" and p and p.type == syms.import_name)
792 or (v == "from" and p and p.type == syms.import_from)
797 def is_with_or_async_with_stmt(leaf: Leaf) -> bool:
798 """Return True if the given leaf starts a with or async with statement."""
800 leaf.type == token.NAME
801 and leaf.value == "with"
803 and leaf.parent.type == syms.with_stmt
805 leaf.type == token.ASYNC
806 and leaf.next_sibling
807 and leaf.next_sibling.type == syms.with_stmt
811 def is_async_stmt_or_funcdef(leaf: Leaf) -> bool:
812 """Return True if the given leaf starts an async def/for/with statement.
814 Note that `async def` can be either an `async_stmt` or `async_funcdef`,
815 the latter is used when it has decorators.
818 leaf.type == token.ASYNC
820 and leaf.parent.type in {syms.async_stmt, syms.async_funcdef}
824 def is_type_comment(leaf: Leaf) -> bool:
825 """Return True if the given leaf is a type comment. This function should only
826 be used for general type comments (excluding ignore annotations, which should
827 use `is_type_ignore_comment`). Note that general type comments are no longer
828 used in modern version of Python, this function may be deprecated in the future."""
831 return t in {token.COMMENT, STANDALONE_COMMENT} and v.startswith("# type:")
834 def is_type_ignore_comment(leaf: Leaf) -> bool:
835 """Return True if the given leaf is a type comment with ignore annotation."""
838 return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_ignore_comment_string(v)
841 def is_type_ignore_comment_string(value: str) -> bool:
842 """Return True if the given string match with type comment with
843 ignore annotation."""
844 return value.startswith("# type: ignore")
847 def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:
848 """Wrap `child` in parentheses.
850 This replaces `child` with an atom holding the parentheses and the old
851 child. That requires moving the prefix.
853 If `visible` is False, the leaves will be valueless (and thus invisible).
855 lpar = Leaf(token.LPAR, "(" if visible else "")
856 rpar = Leaf(token.RPAR, ")" if visible else "")
857 prefix = child.prefix
859 index = child.remove() or 0
860 new_child = Node(syms.atom, [lpar, child, rpar])
861 new_child.prefix = prefix
862 parent.insert_child(index, new_child)
865 def unwrap_singleton_parenthesis(node: LN) -> Optional[LN]:
866 """Returns `wrapped` if `node` is of the shape ( wrapped ).
868 Parenthesis can be optional. Returns None otherwise"""
869 if len(node.children) != 3:
872 lpar, wrapped, rpar = node.children
873 if not (lpar.type == token.LPAR and rpar.type == token.RPAR):
879 def ensure_visible(leaf: Leaf) -> None:
880 """Make sure parentheses are visible.
882 They could be invisible as part of some statements (see
883 :func:`normalize_invisible_parens` and :func:`visit_import_from`).
885 if leaf.type == token.LPAR:
887 elif leaf.type == token.RPAR:
891 def is_name_token(nl: NL) -> TypeGuard[Leaf]:
892 return nl.type == token.NAME
895 def is_lpar_token(nl: NL) -> TypeGuard[Leaf]:
896 return nl.type == token.LPAR
899 def is_rpar_token(nl: NL) -> TypeGuard[Leaf]:
900 return nl.type == token.RPAR
903 def is_string_token(nl: NL) -> TypeGuard[Leaf]:
904 return nl.type == token.STRING
907 def is_number_token(nl: NL) -> TypeGuard[Leaf]:
908 return nl.type == token.NUMBER
911 def is_part_of_annotation(leaf: Leaf) -> bool:
912 """Returns whether this leaf is part of type annotations."""
913 ancestor = leaf.parent
914 while ancestor is not None:
915 if ancestor.prev_sibling and ancestor.prev_sibling.type == token.RARROW:
917 if ancestor.parent and ancestor.parent.type == syms.tname:
919 ancestor = ancestor.parent