| """ |
| blib2to3 Node/Leaf transformation-related utility functions. |
| """ |
|
|
| from collections.abc import Iterator |
| from typing import Final, Generic, Literal, TypeGuard, TypeVar, Union |
|
|
| from mypy_extensions import mypyc_attr |
|
|
| from black.cache import CACHE_DIR |
| from black.mode import Mode, Preview |
| from black.strings import get_string_prefix, has_triple_quotes |
| from blib2to3 import pygram |
| from blib2to3.pgen2 import token |
| from blib2to3.pytree import NL, Leaf, Node, type_repr |
|
|
| pygram.initialize(CACHE_DIR) |
| syms: Final = pygram.python_symbols |
|
|
|
|
| |
| T = TypeVar("T") |
| LN = Union[Leaf, Node] |
| LeafID = int |
| NodeType = int |
|
|
|
|
| WHITESPACE: Final = {token.DEDENT, token.INDENT, token.NEWLINE} |
| STATEMENT: Final = { |
| syms.if_stmt, |
| syms.while_stmt, |
| syms.for_stmt, |
| syms.try_stmt, |
| syms.except_clause, |
| syms.with_stmt, |
| syms.funcdef, |
| syms.classdef, |
| syms.match_stmt, |
| syms.case_block, |
| } |
| STANDALONE_COMMENT: Final = 153 |
| token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT" |
| LOGIC_OPERATORS: Final = {"and", "or"} |
| COMPARATORS: Final = { |
| token.LESS, |
| token.GREATER, |
| token.EQEQUAL, |
| token.NOTEQUAL, |
| token.LESSEQUAL, |
| token.GREATEREQUAL, |
| } |
| MATH_OPERATORS: Final = { |
| token.VBAR, |
| token.CIRCUMFLEX, |
| token.AMPER, |
| token.LEFTSHIFT, |
| token.RIGHTSHIFT, |
| token.PLUS, |
| token.MINUS, |
| token.STAR, |
| token.SLASH, |
| token.DOUBLESLASH, |
| token.PERCENT, |
| token.AT, |
| token.TILDE, |
| token.DOUBLESTAR, |
| } |
| STARS: Final = {token.STAR, token.DOUBLESTAR} |
| VARARGS_SPECIALS: Final = STARS | {token.SLASH} |
| VARARGS_PARENTS: Final = { |
| syms.arglist, |
| syms.argument, |
| syms.trailer, |
| syms.typedargslist, |
| syms.varargslist, |
| } |
| UNPACKING_PARENTS: Final = { |
| syms.atom, |
| syms.dictsetmaker, |
| syms.listmaker, |
| syms.testlist_gexp, |
| syms.testlist_star_expr, |
| syms.subject_expr, |
| syms.pattern, |
| } |
| TEST_DESCENDANTS: Final = { |
| syms.test, |
| syms.lambdef, |
| syms.or_test, |
| syms.and_test, |
| syms.not_test, |
| syms.comparison, |
| syms.star_expr, |
| syms.expr, |
| syms.xor_expr, |
| syms.and_expr, |
| syms.shift_expr, |
| syms.arith_expr, |
| syms.trailer, |
| syms.term, |
| syms.power, |
| syms.namedexpr_test, |
| } |
| TYPED_NAMES: Final = {syms.tname, syms.tname_star} |
| ASSIGNMENTS: Final = { |
| "=", |
| "+=", |
| "-=", |
| "*=", |
| "@=", |
| "/=", |
| "%=", |
| "&=", |
| "|=", |
| "^=", |
| "<<=", |
| ">>=", |
| "**=", |
| "//=", |
| ":", |
| } |
|
|
| IMPLICIT_TUPLE: Final = {syms.testlist, syms.testlist_star_expr, syms.exprlist} |
| BRACKET: Final = { |
| token.LPAR: token.RPAR, |
| token.LSQB: token.RSQB, |
| token.LBRACE: token.RBRACE, |
| } |
| OPENING_BRACKETS: Final = set(BRACKET.keys()) |
| CLOSING_BRACKETS: Final = set(BRACKET.values()) |
| BRACKETS: Final = OPENING_BRACKETS | CLOSING_BRACKETS |
| ALWAYS_NO_SPACE: Final = CLOSING_BRACKETS | { |
| token.COMMA, |
| STANDALONE_COMMENT, |
| token.FSTRING_MIDDLE, |
| token.FSTRING_END, |
| token.TSTRING_MIDDLE, |
| token.TSTRING_END, |
| token.BANG, |
| } |
|
|
| RARROW = 55 |
|
|
|
|
| @mypyc_attr(allow_interpreted_subclasses=True) |
| class Visitor(Generic[T]): |
| """Basic lib2to3 visitor that yields things of type `T` on `visit()`.""" |
|
|
| def visit(self, node: LN) -> Iterator[T]: |
| """Main method to visit `node` and its children. |
| |
| It tries to find a `visit_*()` method for the given `node.type`, like |
| `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects. |
| If no dedicated `visit_*()` method is found, chooses `visit_default()` |
| instead. |
| |
| Then yields objects of type `T` from the selected visitor. |
| """ |
| if node.type < 256: |
| name = token.tok_name[node.type] |
| else: |
| name = str(type_repr(node.type)) |
| |
| |
| |
| |
| visitf = getattr(self, f"visit_{name}", None) |
| if visitf: |
| yield from visitf(node) |
| else: |
| yield from self.visit_default(node) |
|
|
| def visit_default(self, node: LN) -> Iterator[T]: |
| """Default `visit_*()` implementation. Recurses to children of `node`.""" |
| if isinstance(node, Node): |
| for child in node.children: |
| yield from self.visit(child) |
|
|
|
|
| def whitespace(leaf: Leaf, *, complex_subscript: bool, mode: Mode) -> str: |
| """Return whitespace prefix if needed for the given `leaf`. |
| |
| `complex_subscript` signals whether the given leaf is part of a subscription |
| which has non-trivial arguments, like arithmetic expressions or function calls. |
| """ |
| NO: Final[str] = "" |
| SPACE: Final[str] = " " |
| DOUBLESPACE: Final[str] = " " |
| t = leaf.type |
| p = leaf.parent |
| v = leaf.value |
| if t in ALWAYS_NO_SPACE: |
| return NO |
|
|
| if t == token.COMMENT: |
| return DOUBLESPACE |
|
|
| assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}" |
| if t == token.COLON and p.type not in { |
| syms.subscript, |
| syms.subscriptlist, |
| syms.sliceop, |
| }: |
| return NO |
|
|
| if t == token.LBRACE and p.type in ( |
| syms.fstring_replacement_field, |
| syms.tstring_replacement_field, |
| ): |
| return NO |
|
|
| prev = leaf.prev_sibling |
| if not prev: |
| prevp = preceding_leaf(p) |
| if not prevp or prevp.type in OPENING_BRACKETS: |
| return NO |
|
|
| if t == token.COLON: |
| if prevp.type == token.COLON: |
| return NO |
|
|
| elif prevp.type != token.COMMA and not complex_subscript: |
| return NO |
|
|
| return SPACE |
|
|
| if prevp.type == token.EQUAL: |
| if prevp.parent: |
| if prevp.parent.type in { |
| syms.arglist, |
| syms.argument, |
| syms.parameters, |
| syms.varargslist, |
| }: |
| return NO |
|
|
| elif prevp.parent.type == syms.typedargslist: |
| |
| |
| |
| return prevp.prefix |
|
|
| elif ( |
| prevp.type == token.STAR |
| and parent_type(prevp) == syms.star_expr |
| and parent_type(prevp.parent) in (syms.subscriptlist, syms.tname_star) |
| ): |
| |
| return NO |
|
|
| elif prevp.type in VARARGS_SPECIALS: |
| if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS): |
| return NO |
|
|
| elif prevp.type == token.COLON: |
| if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}: |
| return SPACE if complex_subscript else NO |
|
|
| elif ( |
| prevp.parent |
| and prevp.parent.type == syms.factor |
| and prevp.type in MATH_OPERATORS |
| ): |
| return NO |
|
|
| elif prevp.type == token.AT and p.parent and p.parent.type == syms.decorator: |
| |
| return NO |
|
|
| elif prev.type in OPENING_BRACKETS: |
| return NO |
|
|
| elif prev.type == token.BANG: |
| return NO |
|
|
| if p.type in {syms.parameters, syms.arglist}: |
| |
| if not prev or prev.type != token.COMMA: |
| return NO |
|
|
| elif p.type == syms.varargslist: |
| |
| if prev and prev.type != token.COMMA: |
| return NO |
|
|
| elif p.type == syms.typedargslist: |
| |
| if not prev: |
| return NO |
|
|
| if t == token.EQUAL: |
| if prev.type not in TYPED_NAMES: |
| return NO |
|
|
| elif prev.type == token.EQUAL: |
| |
| |
| return prev.prefix |
|
|
| elif prev.type != token.COMMA: |
| return NO |
|
|
| elif p.type in TYPED_NAMES: |
| |
| if not prev: |
| prevp = preceding_leaf(p) |
| if not prevp or prevp.type != token.COMMA: |
| return NO |
|
|
| elif p.type == syms.trailer: |
| |
| if t == token.LPAR or t == token.RPAR: |
| return NO |
|
|
| if not prev: |
| if t == token.DOT or t == token.LSQB: |
| return NO |
|
|
| elif prev.type != token.COMMA: |
| return NO |
|
|
| elif p.type == syms.argument: |
| |
| if t == token.EQUAL: |
| return NO |
|
|
| if not prev: |
| prevp = preceding_leaf(p) |
| if not prevp or prevp.type == token.LPAR: |
| return NO |
|
|
| elif prev.type in {token.EQUAL} | VARARGS_SPECIALS: |
| return NO |
|
|
| elif p.type == syms.decorator: |
| |
| return NO |
|
|
| elif p.type == syms.dotted_name: |
| if prev: |
| return NO |
|
|
| prevp = preceding_leaf(p) |
| if not prevp or prevp.type == token.AT or prevp.type == token.DOT: |
| return NO |
|
|
| elif p.type == syms.classdef: |
| if t == token.LPAR: |
| return NO |
|
|
| if prev and prev.type == token.LPAR: |
| return NO |
|
|
| elif p.type in {syms.subscript, syms.sliceop}: |
| |
| if not prev: |
| assert p.parent is not None, "subscripts are always parented" |
| if p.parent.type == syms.subscriptlist: |
| return SPACE |
|
|
| return NO |
|
|
| elif t == token.COLONEQUAL or prev.type == token.COLONEQUAL: |
| return SPACE |
|
|
| elif not complex_subscript: |
| return NO |
|
|
| elif p.type == syms.atom: |
| if prev and t == token.DOT: |
| |
| return NO |
|
|
| elif p.type == syms.dictsetmaker: |
| |
| if prev and prev.type == token.DOUBLESTAR: |
| return NO |
|
|
| elif p.type in {syms.factor, syms.star_expr}: |
| |
| if not prev: |
| prevp = preceding_leaf(p) |
| if not prevp or prevp.type in OPENING_BRACKETS: |
| return NO |
|
|
| prevp_parent = prevp.parent |
| assert prevp_parent is not None |
| if prevp.type == token.COLON and prevp_parent.type in { |
| syms.subscript, |
| syms.sliceop, |
| }: |
| return NO |
|
|
| elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument: |
| return NO |
|
|
| elif t in {token.NAME, token.NUMBER, token.STRING}: |
| return NO |
|
|
| elif p.type == syms.import_from: |
| if t == token.DOT: |
| if prev and prev.type == token.DOT: |
| return NO |
|
|
| elif t == token.NAME: |
| if v == "import": |
| return SPACE |
|
|
| if prev and prev.type == token.DOT: |
| return NO |
|
|
| elif p.type == syms.sliceop: |
| return NO |
|
|
| elif p.type == syms.except_clause: |
| if t == token.STAR: |
| return NO |
|
|
| return SPACE |
|
|
|
|
| def make_simple_prefix(nl_count: int, form_feed: bool, empty_line: str = "\n") -> str: |
| """Generate a normalized prefix string.""" |
| if form_feed: |
| return (empty_line * (nl_count - 1)) + "\f" + empty_line |
| return empty_line * nl_count |
|
|
|
|
| def preceding_leaf(node: LN | None) -> Leaf | None: |
| """Return the first leaf that precedes `node`, if any.""" |
| while node: |
| res = node.prev_sibling |
| if res: |
| if isinstance(res, Leaf): |
| return res |
|
|
| try: |
| return list(res.leaves())[-1] |
|
|
| except IndexError: |
| return None |
|
|
| node = node.parent |
| return None |
|
|
|
|
| def prev_siblings_are(node: LN | None, tokens: list[NodeType | None]) -> bool: |
| """Return if the `node` and its previous siblings match types against the provided |
| list of tokens; the provided `node`has its type matched against the last element in |
| the list. `None` can be used as the first element to declare that the start of the |
| list is anchored at the start of its parent's children.""" |
| if not tokens: |
| return True |
| if tokens[-1] is None: |
| return node is None |
| if not node: |
| return False |
| if node.type != tokens[-1]: |
| return False |
| return prev_siblings_are(node.prev_sibling, tokens[:-1]) |
|
|
|
|
| def parent_type(node: LN | None) -> NodeType | None: |
| """ |
| Returns: |
| @node.parent.type, if @node is not None and has a parent. |
| OR |
| None, otherwise. |
| """ |
| if node is None or node.parent is None: |
| return None |
|
|
| return node.parent.type |
|
|
|
|
| def child_towards(ancestor: Node, descendant: LN) -> LN | None: |
| """Return the child of `ancestor` that contains `descendant`.""" |
| node: LN | None = descendant |
| while node and node.parent != ancestor: |
| node = node.parent |
| return node |
|
|
|
|
| def replace_child(old_child: LN, new_child: LN) -> None: |
| """ |
| Side Effects: |
| * If @old_child.parent is set, replace @old_child with @new_child in |
| @old_child's underlying Node structure. |
| OR |
| * Otherwise, this function does nothing. |
| """ |
| parent = old_child.parent |
| if not parent: |
| return |
|
|
| child_idx = old_child.remove() |
| if child_idx is not None: |
| parent.insert_child(child_idx, new_child) |
|
|
|
|
| def container_of(leaf: Leaf) -> LN: |
| """Return `leaf` or one of its ancestors that is the topmost container of it. |
| |
| By "container" we mean a node where `leaf` is the very first child. |
| """ |
| same_prefix = leaf.prefix |
| container: LN = leaf |
| while container: |
| parent = container.parent |
| if parent is None: |
| break |
|
|
| if parent.children[0].prefix != same_prefix: |
| break |
|
|
| if parent.type == syms.file_input: |
| break |
|
|
| if parent.prev_sibling is not None and parent.prev_sibling.type in BRACKETS: |
| break |
|
|
| container = parent |
| return container |
|
|
|
|
| def first_leaf_of(node: LN) -> Leaf | None: |
| """Returns the first leaf of the node tree.""" |
| if isinstance(node, Leaf): |
| return node |
| if node.children: |
| return first_leaf_of(node.children[0]) |
| else: |
| return None |
|
|
|
|
| def is_arith_like(node: LN) -> bool: |
| """Whether node is an arithmetic or a binary arithmetic expression""" |
| return node.type in { |
| syms.arith_expr, |
| syms.shift_expr, |
| syms.xor_expr, |
| syms.and_expr, |
| } |
|
|
|
|
| def is_docstring(node: NL) -> bool: |
| if isinstance(node, Leaf): |
| if node.type != token.STRING: |
| return False |
|
|
| prefix = get_string_prefix(node.value) |
| if set(prefix).intersection("bBfF"): |
| return False |
|
|
| if ( |
| node.parent |
| and node.parent.type == syms.simple_stmt |
| and not node.parent.prev_sibling |
| and node.parent.parent |
| and node.parent.parent.type == syms.file_input |
| ): |
| return True |
|
|
| if prev_siblings_are( |
| node.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt] |
| ): |
| return True |
|
|
| |
| if prev_siblings_are(node.parent, [syms.parameters, token.COLON, syms.simple_stmt]): |
| |
| |
| return True |
|
|
| return False |
|
|
|
|
| def is_empty_tuple(node: LN) -> bool: |
| """Return True if `node` holds an empty tuple.""" |
| return ( |
| node.type == syms.atom |
| and len(node.children) == 2 |
| and node.children[0].type == token.LPAR |
| and node.children[1].type == token.RPAR |
| ) |
|
|
|
|
| def is_one_tuple(node: LN) -> bool: |
| """Return True if `node` holds a tuple with one element, with or without parens.""" |
| if node.type == syms.atom: |
| gexp = unwrap_singleton_parenthesis(node) |
| if gexp is None or gexp.type != syms.testlist_gexp: |
| return False |
|
|
| return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA |
|
|
| return ( |
| node.type in IMPLICIT_TUPLE |
| and len(node.children) == 2 |
| and node.children[1].type == token.COMMA |
| ) |
|
|
|
|
| def is_tuple(node: LN) -> bool: |
| """Return True if `node` holds a tuple.""" |
| if node.type != syms.atom: |
| return False |
| gexp = unwrap_singleton_parenthesis(node) |
| if gexp is None or gexp.type != syms.testlist_gexp: |
| return False |
|
|
| return True |
|
|
|
|
| def is_tuple_containing_walrus(node: LN) -> bool: |
| """Return True if `node` holds a tuple that contains a walrus operator.""" |
| if node.type != syms.atom: |
| return False |
| gexp = unwrap_singleton_parenthesis(node) |
| if gexp is None or gexp.type != syms.testlist_gexp: |
| return False |
|
|
| return any(child.type == syms.namedexpr_test for child in gexp.children) |
|
|
|
|
| def is_tuple_containing_star(node: LN) -> bool: |
| """Return True if `node` holds a tuple that contains a star operator.""" |
| if node.type != syms.atom: |
| return False |
| gexp = unwrap_singleton_parenthesis(node) |
| if gexp is None or gexp.type != syms.testlist_gexp: |
| return False |
|
|
| return any(child.type == syms.star_expr for child in gexp.children) |
|
|
|
|
| def is_generator(node: LN) -> bool: |
| """Return True if `node` holds a generator.""" |
| if node.type != syms.atom: |
| return False |
| gexp = unwrap_singleton_parenthesis(node) |
| if gexp is None or gexp.type != syms.testlist_gexp: |
| return False |
|
|
| return any(child.type == syms.old_comp_for for child in gexp.children) |
|
|
|
|
| def is_one_sequence_between( |
| opening: Leaf, |
| closing: Leaf, |
| leaves: list[Leaf], |
| brackets: tuple[int, int] = (token.LPAR, token.RPAR), |
| ) -> bool: |
| """Return True if content between `opening` and `closing` is a one-sequence.""" |
| if (opening.type, closing.type) != brackets: |
| return False |
|
|
| depth = closing.bracket_depth + 1 |
| for _opening_index, leaf in enumerate(leaves): |
| if leaf is opening: |
| break |
|
|
| else: |
| raise LookupError("Opening paren not found in `leaves`") |
|
|
| commas = 0 |
| _opening_index += 1 |
| for leaf in leaves[_opening_index:]: |
| if leaf is closing: |
| break |
|
|
| bracket_depth = leaf.bracket_depth |
| if bracket_depth == depth and leaf.type == token.COMMA: |
| commas += 1 |
| if leaf.parent and leaf.parent.type in { |
| syms.arglist, |
| syms.typedargslist, |
| }: |
| commas += 1 |
| break |
|
|
| return commas < 2 |
|
|
|
|
| def is_walrus_assignment(node: LN) -> bool: |
| """Return True iff `node` is of the shape ( test := test )""" |
| inner = unwrap_singleton_parenthesis(node) |
| return inner is not None and inner.type == syms.namedexpr_test |
|
|
|
|
| def is_simple_decorator_trailer(node: LN, last: bool = False) -> bool: |
| """Return True iff `node` is a trailer valid in a simple decorator""" |
| return node.type == syms.trailer and ( |
| ( |
| len(node.children) == 2 |
| and node.children[0].type == token.DOT |
| and node.children[1].type == token.NAME |
| ) |
| |
| or ( |
| last |
| and len(node.children) == 2 |
| and node.children[0].type == token.LPAR |
| and node.children[1].type == token.RPAR |
| ) |
| |
| or ( |
| last |
| and len(node.children) == 3 |
| and node.children[0].type == token.LPAR |
| |
| and node.children[2].type == token.RPAR |
| ) |
| ) |
|
|
|
|
| def is_simple_decorator_expression(node: LN) -> bool: |
| """Return True iff `node` could be a 'dotted name' decorator |
| |
| This function takes the node of the 'namedexpr_test' of the new decorator |
| grammar and test if it would be valid under the old decorator grammar. |
| |
| The old grammar was: decorator: @ dotted_name [arguments] NEWLINE |
| The new grammar is : decorator: @ namedexpr_test NEWLINE |
| """ |
| if node.type == token.NAME: |
| return True |
| if node.type == syms.power: |
| if node.children: |
| return ( |
| node.children[0].type == token.NAME |
| and all(map(is_simple_decorator_trailer, node.children[1:-1])) |
| and ( |
| len(node.children) < 2 |
| or is_simple_decorator_trailer(node.children[-1], last=True) |
| ) |
| ) |
| return False |
|
|
|
|
| def is_yield(node: LN) -> bool: |
| """Return True if `node` holds a `yield` or `yield from` expression.""" |
| if node.type == syms.yield_expr: |
| return True |
|
|
| if is_name_token(node) and node.value == "yield": |
| return True |
|
|
| if node.type != syms.atom: |
| return False |
|
|
| if len(node.children) != 3: |
| return False |
|
|
| lpar, expr, rpar = node.children |
| if lpar.type == token.LPAR and rpar.type == token.RPAR: |
| return is_yield(expr) |
|
|
| return False |
|
|
|
|
| def is_vararg(leaf: Leaf, within: set[NodeType]) -> bool: |
| """Return True if `leaf` is a star or double star in a vararg or kwarg. |
| |
| If `within` includes VARARGS_PARENTS, this applies to function signatures. |
| If `within` includes UNPACKING_PARENTS, it applies to right hand-side |
| extended iterable unpacking (PEP 3132) and additional unpacking |
| generalizations (PEP 448). |
| """ |
| if leaf.type not in VARARGS_SPECIALS or not leaf.parent: |
| return False |
|
|
| p = leaf.parent |
| if p.type == syms.star_expr: |
| |
| |
| if not p.parent: |
| return False |
|
|
| p = p.parent |
|
|
| return p.type in within |
|
|
|
|
| def is_fstring(node: Node) -> bool: |
| """Return True if the node is an f-string""" |
| return node.type == syms.fstring |
|
|
|
|
| def fstring_tstring_to_string(node: Node) -> Leaf: |
| """Converts an fstring or tstring node back to a string node.""" |
| string_without_prefix = str(node)[len(node.prefix) :] |
| string_leaf = Leaf(token.STRING, string_without_prefix, prefix=node.prefix) |
| string_leaf.lineno = node.get_lineno() or 0 |
| return string_leaf |
|
|
|
|
| def is_multiline_string(node: LN) -> bool: |
| """Return True if `leaf` is a multiline string that actually spans many lines.""" |
| if isinstance(node, Node) and is_fstring(node): |
| leaf = fstring_tstring_to_string(node) |
| elif isinstance(node, Leaf): |
| leaf = node |
| else: |
| return False |
|
|
| return has_triple_quotes(leaf.value) and "\n" in leaf.value |
|
|
|
|
| def is_parent_function_or_class(node: Node) -> bool: |
| assert node.type in {syms.suite, syms.simple_stmt} |
| assert node.parent is not None |
| |
| return node.parent.type in {syms.funcdef, syms.classdef} |
|
|
|
|
| def is_function_or_class(node: Node) -> bool: |
| return node.type in {syms.funcdef, syms.classdef, syms.async_funcdef} |
|
|
|
|
| def is_stub_suite(node: Node) -> bool: |
| """Return True if `node` is a suite with a stub body.""" |
| if node.parent is not None and not is_parent_function_or_class(node): |
| return False |
|
|
| |
| if node.prefix.strip(): |
| return False |
|
|
| if ( |
| len(node.children) != 4 |
| or node.children[0].type != token.NEWLINE |
| or node.children[1].type != token.INDENT |
| or node.children[3].type != token.DEDENT |
| ): |
| return False |
|
|
| if node.children[3].prefix.strip(): |
| return False |
|
|
| return is_stub_body(node.children[2]) |
|
|
|
|
| def is_stub_body(node: LN) -> bool: |
| """Return True if `node` is a simple statement containing an ellipsis.""" |
| if not isinstance(node, Node) or node.type != syms.simple_stmt: |
| return False |
|
|
| if len(node.children) != 2: |
| return False |
|
|
| child = node.children[0] |
| return ( |
| not child.prefix.strip() |
| and child.type == syms.atom |
| and len(child.children) == 3 |
| and all(leaf == Leaf(token.DOT, ".") for leaf in child.children) |
| ) |
|
|
|
|
| def is_atom_with_invisible_parens(node: LN) -> bool: |
| """Given a `LN`, determines whether it's an atom `node` with invisible |
| parens. Useful in dedupe-ing and normalizing parens. |
| """ |
| if isinstance(node, Leaf) or node.type != syms.atom: |
| return False |
|
|
| first, last = node.children[0], node.children[-1] |
| return ( |
| isinstance(first, Leaf) |
| and first.type == token.LPAR |
| and first.value == "" |
| and isinstance(last, Leaf) |
| and last.type == token.RPAR |
| and last.value == "" |
| ) |
|
|
|
|
| def is_empty_par(leaf: Leaf) -> bool: |
| return is_empty_lpar(leaf) or is_empty_rpar(leaf) |
|
|
|
|
| def is_empty_lpar(leaf: Leaf) -> bool: |
| return leaf.type == token.LPAR and leaf.value == "" |
|
|
|
|
| def is_empty_rpar(leaf: Leaf) -> bool: |
| return leaf.type == token.RPAR and leaf.value == "" |
|
|
|
|
| def is_import(leaf: Leaf) -> bool: |
| """Return True if the given leaf starts an import statement.""" |
| p = leaf.parent |
| t = leaf.type |
| v = leaf.value |
| return bool( |
| t == token.NAME |
| and ( |
| (v == "import" and p and p.type == syms.import_name) |
| or (v == "from" and p and p.type == syms.import_from) |
| ) |
| ) |
|
|
|
|
| def is_with_or_async_with_stmt(leaf: Leaf) -> bool: |
| """Return True if the given leaf starts a with or async with statement.""" |
| return bool( |
| leaf.type == token.NAME |
| and leaf.value == "with" |
| and leaf.parent |
| and leaf.parent.type == syms.with_stmt |
| ) or bool( |
| leaf.type == token.ASYNC |
| and leaf.next_sibling |
| and leaf.next_sibling.type == syms.with_stmt |
| ) |
|
|
|
|
| def is_async_stmt_or_funcdef(leaf: Leaf) -> bool: |
| """Return True if the given leaf starts an async def/for/with statement. |
| |
| Note that `async def` can be either an `async_stmt` or `async_funcdef`, |
| the latter is used when it has decorators. |
| """ |
| return bool( |
| leaf.type == token.ASYNC |
| and leaf.parent |
| and leaf.parent.type in {syms.async_stmt, syms.async_funcdef} |
| ) |
|
|
|
|
| def is_type_comment(leaf: Leaf, mode: Mode) -> bool: |
| """Return True if the given leaf is a type comment. This function should only |
| be used for general type comments (excluding ignore annotations, which should |
| use `is_type_ignore_comment`). Note that general type comments are no longer |
| used in modern version of Python, this function may be deprecated in the future.""" |
| t = leaf.type |
| v = leaf.value |
| return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_comment_string(v, mode) |
|
|
|
|
| def is_type_comment_string(value: str, mode: Mode) -> bool: |
| if Preview.standardize_type_comments in mode: |
| is_valid = value.startswith("#") and value[1:].lstrip().startswith("type:") |
| else: |
| is_valid = value.startswith("# type:") |
| return is_valid |
|
|
|
|
| def is_type_ignore_comment(leaf: Leaf, mode: Mode) -> bool: |
| """Return True if the given leaf is a type comment with ignore annotation.""" |
| t = leaf.type |
| v = leaf.value |
| return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_ignore_comment_string( |
| v, mode |
| ) |
|
|
|
|
| def is_type_ignore_comment_string(value: str, mode: Mode) -> bool: |
| """Return True if the given string match with type comment with |
| ignore annotation.""" |
| if Preview.standardize_type_comments in mode: |
| is_valid = is_type_comment_string(value, mode) and value.split(":", 1)[ |
| 1 |
| ].lstrip().startswith("ignore") |
| else: |
| is_valid = value.startswith("# type: ignore") |
|
|
| return is_valid |
|
|
|
|
| def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None: |
| """Wrap `child` in parentheses. |
| |
| This replaces `child` with an atom holding the parentheses and the old |
| child. That requires moving the prefix. |
| |
| If `visible` is False, the leaves will be valueless (and thus invisible). |
| """ |
| lpar = Leaf(token.LPAR, "(" if visible else "") |
| rpar = Leaf(token.RPAR, ")" if visible else "") |
| prefix = child.prefix |
| child.prefix = "" |
| index = child.remove() or 0 |
| new_child = Node(syms.atom, [lpar, child, rpar]) |
| new_child.prefix = prefix |
| parent.insert_child(index, new_child) |
|
|
|
|
| def unwrap_singleton_parenthesis(node: LN) -> LN | None: |
| """Returns `wrapped` if `node` is of the shape ( wrapped ). |
| |
| Parenthesis can be optional. Returns None otherwise""" |
| if len(node.children) != 3: |
| return None |
|
|
| lpar, wrapped, rpar = node.children |
| if not (lpar.type == token.LPAR and rpar.type == token.RPAR): |
| return None |
|
|
| return wrapped |
|
|
|
|
| def ensure_visible(leaf: Leaf) -> None: |
| """Make sure parentheses are visible. |
| |
| They could be invisible as part of some statements (see |
| :func:`normalize_invisible_parens` and :func:`visit_import_from`). |
| """ |
| if leaf.type == token.LPAR: |
| leaf.value = "(" |
| elif leaf.type == token.RPAR: |
| leaf.value = ")" |
|
|
|
|
| def is_name_token(nl: NL) -> TypeGuard[Leaf]: |
| return nl.type == token.NAME |
|
|
|
|
| def is_lpar_token(nl: NL) -> TypeGuard[Leaf]: |
| return nl.type == token.LPAR |
|
|
|
|
| def is_rpar_token(nl: NL) -> TypeGuard[Leaf]: |
| return nl.type == token.RPAR |
|
|
|
|
| def is_number_token(nl: NL) -> TypeGuard[Leaf]: |
| return nl.type == token.NUMBER |
|
|
|
|
| def get_annotation_type(leaf: Leaf) -> Literal["return", "param", None]: |
| """Returns the type of annotation this leaf is part of, if any.""" |
| ancestor = leaf.parent |
| while ancestor is not None: |
| if ancestor.prev_sibling and ancestor.prev_sibling.type == token.RARROW: |
| return "return" |
| if ancestor.parent and ancestor.parent.type == syms.tname: |
| return "param" |
| ancestor = ancestor.parent |
| return None |
|
|
|
|
| def is_part_of_annotation(leaf: Leaf) -> bool: |
| """Returns whether this leaf is part of a type annotation.""" |
| assert leaf.parent is not None |
| return get_annotation_type(leaf) is not None |
|
|
|
|
| def first_leaf(node: LN) -> Leaf | None: |
| """Returns the first leaf of the ancestor node.""" |
| if isinstance(node, Leaf): |
| return node |
| elif not node.children: |
| return None |
| else: |
| return first_leaf(node.children[0]) |
|
|
|
|
| def last_leaf(node: LN) -> Leaf | None: |
| """Returns the last leaf of the ancestor node.""" |
| if isinstance(node, Leaf): |
| return node |
| elif not node.children: |
| return None |
| else: |
| return last_leaf(node.children[-1]) |
|
|
|
|
| def furthest_ancestor_with_last_leaf(leaf: Leaf) -> LN: |
| """Returns the furthest ancestor that has this leaf node as the last leaf.""" |
| node: LN = leaf |
| while node.parent and node.parent.children and node is node.parent.children[-1]: |
| node = node.parent |
| return node |
|
|
|
|
| def has_sibling_with_type(node: LN, type: int) -> bool: |
| |
| sibling = node.prev_sibling |
| while sibling is not None: |
| if sibling.type == type: |
| return True |
| sibling = sibling.prev_sibling |
|
|
| |
| sibling = node.next_sibling |
| while sibling is not None: |
| if sibling.type == type: |
| return True |
| sibling = sibling.next_sibling |
|
|
| return False |
|
|