Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/black/nodes.py: 19%

1"""

2blib2to3 Node/Leaf transformation-related utility functions.

3"""

5import sys

6from collections.abc import Iterator

7from typing import Final, Generic, Literal, Optional, TypeVar, Union

9if sys.version_info >= (3, 10):

10 from typing import TypeGuard

11else:

12 from typing_extensions import TypeGuard

14from mypy_extensions import mypyc_attr

16from black.cache import CACHE_DIR

17from black.mode import Mode, Preview

18from black.strings import get_string_prefix, has_triple_quotes

19from blib2to3 import pygram

20from blib2to3.pgen2 import token

21from blib2to3.pytree import NL, Leaf, Node, type_repr

23pygram.initialize(CACHE_DIR)

24syms: Final = pygram.python_symbols

27# types

28T = TypeVar("T")

29LN = Union[Leaf, Node]

30LeafID = int

31NodeType = int

34WHITESPACE: Final = {token.DEDENT, token.INDENT, token.NEWLINE}

35STATEMENT: Final = {

36 syms.if_stmt,

37 syms.while_stmt,

38 syms.for_stmt,

39 syms.try_stmt,

40 syms.except_clause,

41 syms.with_stmt,

42 syms.funcdef,

43 syms.classdef,

44 syms.match_stmt,

45 syms.case_block,

46}

47STANDALONE_COMMENT: Final = 153

48token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"

49LOGIC_OPERATORS: Final = {"and", "or"}

50COMPARATORS: Final = {

51 token.LESS,

52 token.GREATER,

53 token.EQEQUAL,

54 token.NOTEQUAL,

55 token.LESSEQUAL,

56 token.GREATEREQUAL,

57}

58MATH_OPERATORS: Final = {

59 token.VBAR,

60 token.CIRCUMFLEX,

61 token.AMPER,

62 token.LEFTSHIFT,

63 token.RIGHTSHIFT,

64 token.PLUS,

65 token.MINUS,

66 token.STAR,

67 token.SLASH,

68 token.DOUBLESLASH,

69 token.PERCENT,

70 token.AT,

71 token.TILDE,

72 token.DOUBLESTAR,

73}

74STARS: Final = {token.STAR, token.DOUBLESTAR}

75VARARGS_SPECIALS: Final = STARS | {token.SLASH}

76VARARGS_PARENTS: Final = {

77 syms.arglist,

78 syms.argument, # double star in arglist

79 syms.trailer, # single argument to call

80 syms.typedargslist,

81 syms.varargslist, # lambdas

82}

83UNPACKING_PARENTS: Final = {

84 syms.atom, # single element of a list or set literal

85 syms.dictsetmaker,

86 syms.listmaker,

87 syms.testlist_gexp,

88 syms.testlist_star_expr,

89 syms.subject_expr,

90 syms.pattern,

91}

92TEST_DESCENDANTS: Final = {

93 syms.test,

94 syms.lambdef,

95 syms.or_test,

96 syms.and_test,

97 syms.not_test,

98 syms.comparison,

99 syms.star_expr,

100 syms.expr,

101 syms.xor_expr,

102 syms.and_expr,

103 syms.shift_expr,

104 syms.arith_expr,

105 syms.trailer,

106 syms.term,

107 syms.power,

108 syms.namedexpr_test,

109}

110TYPED_NAMES: Final = {syms.tname, syms.tname_star}

111ASSIGNMENTS: Final = {

112 "=",

113 "+=",

114 "-=",

115 "*=",

116 "@=",

117 "/=",

118 "%=",

119 "&=",

120 "|=",

121 "^=",

122 "<<=",

123 ">>=",

124 "**=",

125 "//=",

126 ":",

127}

128

129IMPLICIT_TUPLE: Final = {syms.testlist, syms.testlist_star_expr, syms.exprlist}

130BRACKET: Final = {

131 token.LPAR: token.RPAR,

132 token.LSQB: token.RSQB,

133 token.LBRACE: token.RBRACE,

134}

135OPENING_BRACKETS: Final = set(BRACKET.keys())

136CLOSING_BRACKETS: Final = set(BRACKET.values())

137BRACKETS: Final = OPENING_BRACKETS | CLOSING_BRACKETS

138ALWAYS_NO_SPACE: Final = CLOSING_BRACKETS | {

139 token.COMMA,

140 STANDALONE_COMMENT,

141 token.FSTRING_MIDDLE,

142 token.FSTRING_END,

143 token.TSTRING_MIDDLE,

144 token.TSTRING_END,

145 token.BANG,

146}

147

148RARROW = 55

149

150

151@mypyc_attr(allow_interpreted_subclasses=True)

152class Visitor(Generic[T]):

153 """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""

154

155 def visit(self, node: LN) -> Iterator[T]:

156 """Main method to visit `node` and its children.

157

158 It tries to find a `visit_*()` method for the given `node.type`, like

159 `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.

160 If no dedicated `visit_*()` method is found, chooses `visit_default()`

161 instead.

162

163 Then yields objects of type `T` from the selected visitor.

164 """

165 if node.type < 256:

166 name = token.tok_name[node.type]

167 else:

168 name = str(type_repr(node.type))

169 # We explicitly branch on whether a visitor exists (instead of

170 # using self.visit_default as the default arg to getattr) in order

171 # to save needing to create a bound method object and so mypyc can

172 # generate a native call to visit_default.

173 visitf = getattr(self, f"visit_{name}", None)

174 if visitf:

175 yield from visitf(node)

176 else:

177 yield from self.visit_default(node)

178

179 def visit_default(self, node: LN) -> Iterator[T]:

180 """Default `visit_*()` implementation. Recurses to children of `node`."""

181 if isinstance(node, Node):

182 for child in node.children:

183 yield from self.visit(child)

184

185

186def whitespace(leaf: Leaf, *, complex_subscript: bool, mode: Mode) -> str: # noqa: C901

187 """Return whitespace prefix if needed for the given `leaf`.

188

189 `complex_subscript` signals whether the given leaf is part of a subscription

190 which has non-trivial arguments, like arithmetic expressions or function calls.

191 """

192 NO: Final[str] = ""

193 SPACE: Final[str] = " "

194 DOUBLESPACE: Final[str] = " "

195 t = leaf.type

196 p = leaf.parent

197 v = leaf.value

198 if t in ALWAYS_NO_SPACE:

199 return NO

200

201 if t == token.COMMENT:

202 return DOUBLESPACE

203

204 assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"

205 if t == token.COLON and p.type not in {

206 syms.subscript,

207 syms.subscriptlist,

208 syms.sliceop,

209 }:

210 return NO

211

212 if t == token.LBRACE and p.type in (

213 syms.fstring_replacement_field,

214 syms.tstring_replacement_field,

215 ):

216 return NO

217

218 prev = leaf.prev_sibling

219 if not prev:

220 prevp = preceding_leaf(p)

221 if not prevp or prevp.type in OPENING_BRACKETS:

222 return NO

223

224 if t == token.COLON:

225 if prevp.type == token.COLON:

226 return NO

227

228 elif prevp.type != token.COMMA and not complex_subscript:

229 return NO

230

231 return SPACE

232

233 if prevp.type == token.EQUAL:

234 if prevp.parent:

235 if prevp.parent.type in {

236 syms.arglist,

237 syms.argument,

238 syms.parameters,

239 syms.varargslist,

240 }:

241 return NO

242

243 elif prevp.parent.type == syms.typedargslist:

244 # A bit hacky: if the equal sign has whitespace, it means we

245 # previously found it's a typed argument. So, we're using

246 # that, too.

247 return prevp.prefix

248

249 elif (

250 prevp.type == token.STAR

251 and parent_type(prevp) == syms.star_expr

252 and parent_type(prevp.parent) in (syms.subscriptlist, syms.tname_star)

253 ):

254 # No space between typevar tuples or unpacking them.

255 return NO

256

257 elif prevp.type in VARARGS_SPECIALS:

258 if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):

259 return NO

260

261 elif prevp.type == token.COLON:

262 if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:

263 return SPACE if complex_subscript else NO

264

265 elif (

266 prevp.parent

267 and prevp.parent.type == syms.factor

268 and prevp.type in MATH_OPERATORS

269 ):

270 return NO

271

272 elif prevp.type == token.AT and p.parent and p.parent.type == syms.decorator:

273 # no space in decorators

274 return NO

275

276 elif prev.type in OPENING_BRACKETS:

277 return NO

278

279 elif prev.type == token.BANG:

280 return NO

281

282 if p.type in {syms.parameters, syms.arglist}:

283 # untyped function signatures or calls

284 if not prev or prev.type != token.COMMA:

285 return NO

286

287 elif p.type == syms.varargslist:

288 # lambdas

289 if prev and prev.type != token.COMMA:

290 return NO

291

292 elif p.type == syms.typedargslist:

293 # typed function signatures

294 if not prev:

295 return NO

296

297 if t == token.EQUAL:

298 if prev.type not in TYPED_NAMES:

299 return NO

300

301 elif prev.type == token.EQUAL:

302 # A bit hacky: if the equal sign has whitespace, it means we

303 # previously found it's a typed argument. So, we're using that, too.

304 return prev.prefix

305

306 elif prev.type != token.COMMA:

307 return NO

308

309 elif p.type in TYPED_NAMES:

310 # type names

311 if not prev:

312 prevp = preceding_leaf(p)

313 if not prevp or prevp.type != token.COMMA:

314 return NO

315

316 elif p.type == syms.trailer:

317 # attributes and calls

318 if t == token.LPAR or t == token.RPAR:

319 return NO

320

321 if not prev:

322 if t == token.DOT or t == token.LSQB:

323 return NO

324

325 elif prev.type != token.COMMA:

326 return NO

327

328 elif p.type == syms.argument:

329 # single argument

330 if t == token.EQUAL:

331 return NO

332

333 if not prev:

334 prevp = preceding_leaf(p)

335 if not prevp or prevp.type == token.LPAR:

336 return NO

337

338 elif prev.type in {token.EQUAL} | VARARGS_SPECIALS:

339 return NO

340

341 elif p.type == syms.decorator:

342 # decorators

343 return NO

344

345 elif p.type == syms.dotted_name:

346 if prev:

347 return NO

348

349 prevp = preceding_leaf(p)

350 if not prevp or prevp.type == token.AT or prevp.type == token.DOT:

351 return NO

352

353 elif p.type == syms.classdef:

354 if t == token.LPAR:

355 return NO

356

357 if prev and prev.type == token.LPAR:

358 return NO

359

360 elif p.type in {syms.subscript, syms.sliceop}:

361 # indexing

362 if not prev:

363 assert p.parent is not None, "subscripts are always parented"

364 if p.parent.type == syms.subscriptlist:

365 return SPACE

366

367 return NO

368

369 elif t == token.COLONEQUAL or prev.type == token.COLONEQUAL:

370 return SPACE

371

372 elif not complex_subscript:

373 return NO

374

375 elif p.type == syms.atom:

376 if prev and t == token.DOT:

377 # dots, but not the first one.

378 return NO

379

380 elif p.type == syms.dictsetmaker:

381 # dict unpacking

382 if prev and prev.type == token.DOUBLESTAR:

383 return NO

384

385 elif p.type in {syms.factor, syms.star_expr}:

386 # unary ops

387 if not prev:

388 prevp = preceding_leaf(p)

389 if not prevp or prevp.type in OPENING_BRACKETS:

390 return NO

391

392 prevp_parent = prevp.parent

393 assert prevp_parent is not None

394 if prevp.type == token.COLON and prevp_parent.type in {

395 syms.subscript,

396 syms.sliceop,

397 }:

398 return NO

399

400 elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:

401 return NO

402

403 elif t in {token.NAME, token.NUMBER, token.STRING}:

404 return NO

405

406 elif p.type == syms.import_from:

407 if t == token.DOT:

408 if prev and prev.type == token.DOT:

409 return NO

410

411 elif t == token.NAME:

412 if v == "import":

413 return SPACE

414

415 if prev and prev.type == token.DOT:

416 return NO

417

418 elif p.type == syms.sliceop:

419 return NO

420

421 elif p.type == syms.except_clause:

422 if t == token.STAR:

423 return NO

424

425 return SPACE

426

427

428def make_simple_prefix(nl_count: int, form_feed: bool, empty_line: str = "\n") -> str:

429 """Generate a normalized prefix string."""

430 if form_feed:

431 return (empty_line * (nl_count - 1)) + "\f" + empty_line

432 return empty_line * nl_count

433

434

435def preceding_leaf(node: Optional[LN]) -> Optional[Leaf]:

436 """Return the first leaf that precedes `node`, if any."""

437 while node:

438 res = node.prev_sibling

439 if res:

440 if isinstance(res, Leaf):

441 return res

442

443 try:

444 return list(res.leaves())[-1]

445

446 except IndexError:

447 return None

448

449 node = node.parent

450 return None

451

452

453def prev_siblings_are(node: Optional[LN], tokens: list[Optional[NodeType]]) -> bool:

454 """Return if the `node` and its previous siblings match types against the provided

455 list of tokens; the provided `node`has its type matched against the last element in

456 the list. `None` can be used as the first element to declare that the start of the

457 list is anchored at the start of its parent's children."""

458 if not tokens:

459 return True

460 if tokens[-1] is None:

461 return node is None

462 if not node:

463 return False

464 if node.type != tokens[-1]:

465 return False

466 return prev_siblings_are(node.prev_sibling, tokens[:-1])

467

468

469def parent_type(node: Optional[LN]) -> Optional[NodeType]:

470 """

471 Returns:

472 @node.parent.type, if @node is not None and has a parent.

473 OR

474 None, otherwise.

475 """

476 if node is None or node.parent is None:

477 return None

478

479 return node.parent.type

480

481

482def child_towards(ancestor: Node, descendant: LN) -> Optional[LN]:

483 """Return the child of `ancestor` that contains `descendant`."""

484 node: Optional[LN] = descendant

485 while node and node.parent != ancestor:

486 node = node.parent

487 return node

488

489

490def replace_child(old_child: LN, new_child: LN) -> None:

491 """

492 Side Effects:

493 * If @old_child.parent is set, replace @old_child with @new_child in

494 @old_child's underlying Node structure.

495 OR

496 * Otherwise, this function does nothing.

497 """

498 parent = old_child.parent

499 if not parent:

500 return

501

502 child_idx = old_child.remove()

503 if child_idx is not None:

504 parent.insert_child(child_idx, new_child)

505

506

507def container_of(leaf: Leaf) -> LN:

508 """Return `leaf` or one of its ancestors that is the topmost container of it.

509

510 By "container" we mean a node where `leaf` is the very first child.

511 """

512 same_prefix = leaf.prefix

513 container: LN = leaf

514 while container:

515 parent = container.parent

516 if parent is None:

517 break

518

519 if parent.children[0].prefix != same_prefix:

520 break

521

522 if parent.type == syms.file_input:

523 break

524

525 if parent.prev_sibling is not None and parent.prev_sibling.type in BRACKETS:

526 break

527

528 container = parent

529 return container

530

531

532def first_leaf_of(node: LN) -> Optional[Leaf]:

533 """Returns the first leaf of the node tree."""

534 if isinstance(node, Leaf):

535 return node

536 if node.children:

537 return first_leaf_of(node.children[0])

538 else:

539 return None

540

541

542def is_arith_like(node: LN) -> bool:

543 """Whether node is an arithmetic or a binary arithmetic expression"""

544 return node.type in {

545 syms.arith_expr,

546 syms.shift_expr,

547 syms.xor_expr,

548 syms.and_expr,

549 }

550

551

552def is_docstring(node: NL) -> bool:

553 if isinstance(node, Leaf):

554 if node.type != token.STRING:

555 return False

556

557 prefix = get_string_prefix(node.value)

558 if set(prefix).intersection("bBfF"):

559 return False

560

561 if (

562 node.parent

563 and node.parent.type == syms.simple_stmt

564 and not node.parent.prev_sibling

565 and node.parent.parent

566 and node.parent.parent.type == syms.file_input

567 ):

568 return True

569

570 if prev_siblings_are(

571 node.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]

572 ):

573 return True

574

575 # Multiline docstring on the same line as the `def`.

576 if prev_siblings_are(node.parent, [syms.parameters, token.COLON, syms.simple_stmt]):

577 # `syms.parameters` is only used in funcdefs and async_funcdefs in the Python

578 # grammar. We're safe to return True without further checks.

579 return True

580

581 return False

582

583

584def is_empty_tuple(node: LN) -> bool:

585 """Return True if `node` holds an empty tuple."""

586 return (

587 node.type == syms.atom

588 and len(node.children) == 2

589 and node.children[0].type == token.LPAR

590 and node.children[1].type == token.RPAR

591 )

592

593

594def is_one_tuple(node: LN) -> bool:

595 """Return True if `node` holds a tuple with one element, with or without parens."""

596 if node.type == syms.atom:

597 gexp = unwrap_singleton_parenthesis(node)

598 if gexp is None or gexp.type != syms.testlist_gexp:

599 return False

600

601 return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA

602

603 return (

604 node.type in IMPLICIT_TUPLE

605 and len(node.children) == 2

606 and node.children[1].type == token.COMMA

607 )

608

609

610def is_tuple(node: LN) -> bool:

611 """Return True if `node` holds a tuple."""

612 if node.type != syms.atom:

613 return False

614 gexp = unwrap_singleton_parenthesis(node)

615 if gexp is None or gexp.type != syms.testlist_gexp:

616 return False

617

618 return True

619

620

621def is_tuple_containing_walrus(node: LN) -> bool:

622 """Return True if `node` holds a tuple that contains a walrus operator."""

623 if node.type != syms.atom:

624 return False

625 gexp = unwrap_singleton_parenthesis(node)

626 if gexp is None or gexp.type != syms.testlist_gexp:

627 return False

628

629 return any(child.type == syms.namedexpr_test for child in gexp.children)

630

631

632def is_tuple_containing_star(node: LN) -> bool:

633 """Return True if `node` holds a tuple that contains a star operator."""

634 if node.type != syms.atom:

635 return False

636 gexp = unwrap_singleton_parenthesis(node)

637 if gexp is None or gexp.type != syms.testlist_gexp:

638 return False

639

640 return any(child.type == syms.star_expr for child in gexp.children)

641

642

643def is_generator(node: LN) -> bool:

644 """Return True if `node` holds a generator."""

645 if node.type != syms.atom:

646 return False

647 gexp = unwrap_singleton_parenthesis(node)

648 if gexp is None or gexp.type != syms.testlist_gexp:

649 return False

650

651 return any(child.type == syms.old_comp_for for child in gexp.children)

652

653

654def is_one_sequence_between(

655 opening: Leaf,

656 closing: Leaf,

657 leaves: list[Leaf],

658 brackets: tuple[int, int] = (token.LPAR, token.RPAR),

659) -> bool:

660 """Return True if content between `opening` and `closing` is a one-sequence."""

661 if (opening.type, closing.type) != brackets:

662 return False

663

664 depth = closing.bracket_depth + 1

665 for _opening_index, leaf in enumerate(leaves):

666 if leaf is opening:

667 break

668

669 else:

670 raise LookupError("Opening paren not found in `leaves`")

671

672 commas = 0

673 _opening_index += 1

674 for leaf in leaves[_opening_index:]:

675 if leaf is closing:

676 break

677

678 bracket_depth = leaf.bracket_depth

679 if bracket_depth == depth and leaf.type == token.COMMA:

680 commas += 1

681 if leaf.parent and leaf.parent.type in {

682 syms.arglist,

683 syms.typedargslist,

684 }:

685 commas += 1

686 break

687

688 return commas < 2

689

690

691def is_walrus_assignment(node: LN) -> bool:

692 """Return True iff `node` is of the shape ( test := test )"""

693 inner = unwrap_singleton_parenthesis(node)

694 return inner is not None and inner.type == syms.namedexpr_test

695

696

697def is_simple_decorator_trailer(node: LN, last: bool = False) -> bool:

698 """Return True iff `node` is a trailer valid in a simple decorator"""

699 return node.type == syms.trailer and (

700 (

701 len(node.children) == 2

702 and node.children[0].type == token.DOT

703 and node.children[1].type == token.NAME

704 )

705 # last trailer can be an argument-less parentheses pair

706 or (

707 last

708 and len(node.children) == 2

709 and node.children[0].type == token.LPAR

710 and node.children[1].type == token.RPAR

711 )

712 # last trailer can be arguments

713 or (

714 last

715 and len(node.children) == 3

716 and node.children[0].type == token.LPAR

717 # and node.children[1].type == syms.argument

718 and node.children[2].type == token.RPAR

719 )

720 )

721

722

723def is_simple_decorator_expression(node: LN) -> bool:

724 """Return True iff `node` could be a 'dotted name' decorator

725

726 This function takes the node of the 'namedexpr_test' of the new decorator

727 grammar and test if it would be valid under the old decorator grammar.

728

729 The old grammar was: decorator: @ dotted_name [arguments] NEWLINE

730 The new grammar is : decorator: @ namedexpr_test NEWLINE

731 """

732 if node.type == token.NAME:

733 return True

734 if node.type == syms.power:

735 if node.children:

736 return (

737 node.children[0].type == token.NAME

738 and all(map(is_simple_decorator_trailer, node.children[1:-1]))

739 and (

740 len(node.children) < 2

741 or is_simple_decorator_trailer(node.children[-1], last=True)

742 )

743 )

744 return False

745

746

747def is_yield(node: LN) -> bool:

748 """Return True if `node` holds a `yield` or `yield from` expression."""

749 if node.type == syms.yield_expr:

750 return True

751

752 if is_name_token(node) and node.value == "yield":

753 return True

754

755 if node.type != syms.atom:

756 return False

757

758 if len(node.children) != 3:

759 return False

760

761 lpar, expr, rpar = node.children

762 if lpar.type == token.LPAR and rpar.type == token.RPAR:

763 return is_yield(expr)

764

765 return False

766

767

768def is_vararg(leaf: Leaf, within: set[NodeType]) -> bool:

769 """Return True if `leaf` is a star or double star in a vararg or kwarg.

770

771 If `within` includes VARARGS_PARENTS, this applies to function signatures.

772 If `within` includes UNPACKING_PARENTS, it applies to right hand-side

773 extended iterable unpacking (PEP 3132) and additional unpacking

774 generalizations (PEP 448).

775 """

776 if leaf.type not in VARARGS_SPECIALS or not leaf.parent:

777 return False

778

779 p = leaf.parent

780 if p.type == syms.star_expr:

781 # Star expressions are also used as assignment targets in extended

782 # iterable unpacking (PEP 3132). See what its parent is instead.

783 if not p.parent:

784 return False

785

786 p = p.parent

787

788 return p.type in within

789

790

791def is_fstring(node: Node) -> bool:

792 """Return True if the node is an f-string"""

793 return node.type == syms.fstring

794

795

796def fstring_tstring_to_string(node: Node) -> Leaf:

797 """Converts an fstring or tstring node back to a string node."""

798 string_without_prefix = str(node)[len(node.prefix) :]

799 string_leaf = Leaf(token.STRING, string_without_prefix, prefix=node.prefix)

800 string_leaf.lineno = node.get_lineno() or 0

801 return string_leaf

802

803

804def is_multiline_string(node: LN) -> bool:

805 """Return True if `leaf` is a multiline string that actually spans many lines."""

806 if isinstance(node, Node) and is_fstring(node):

807 leaf = fstring_tstring_to_string(node)

808 elif isinstance(node, Leaf):

809 leaf = node

810 else:

811 return False

812

813 return has_triple_quotes(leaf.value) and "\n" in leaf.value

814

815

816def is_parent_function_or_class(node: Node) -> bool:

817 assert node.type in {syms.suite, syms.simple_stmt}

818 assert node.parent is not None

819 # Note this works for suites / simple_stmts in async def as well

820 return node.parent.type in {syms.funcdef, syms.classdef}

821

822

823def is_function_or_class(node: Node) -> bool:

824 return node.type in {syms.funcdef, syms.classdef, syms.async_funcdef}

825

826

827def is_stub_suite(node: Node) -> bool:

828 """Return True if `node` is a suite with a stub body."""

829 if node.parent is not None and not is_parent_function_or_class(node):

830 return False

831

832 # If there is a comment, we want to keep it.

833 if node.prefix.strip():

834 return False

835

836 if (

837 len(node.children) != 4

838 or node.children[0].type != token.NEWLINE

839 or node.children[1].type != token.INDENT

840 or node.children[3].type != token.DEDENT

841 ):

842 return False

843

844 if node.children[3].prefix.strip():

845 return False

846

847 return is_stub_body(node.children[2])

848

849

850def is_stub_body(node: LN) -> bool:

851 """Return True if `node` is a simple statement containing an ellipsis."""

852 if not isinstance(node, Node) or node.type != syms.simple_stmt:

853 return False

854

855 if len(node.children) != 2:

856 return False

857

858 child = node.children[0]

859 return (

860 not child.prefix.strip()

861 and child.type == syms.atom

862 and len(child.children) == 3

863 and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)

864 )

865

866

867def is_atom_with_invisible_parens(node: LN) -> bool:

868 """Given a `LN`, determines whether it's an atom `node` with invisible

869 parens. Useful in dedupe-ing and normalizing parens.

870 """

871 if isinstance(node, Leaf) or node.type != syms.atom:

872 return False

873

874 first, last = node.children[0], node.children[-1]

875 return (

876 isinstance(first, Leaf)

877 and first.type == token.LPAR

878 and first.value == ""

879 and isinstance(last, Leaf)

880 and last.type == token.RPAR

881 and last.value == ""

882 )

883

884

885def is_empty_par(leaf: Leaf) -> bool:

886 return is_empty_lpar(leaf) or is_empty_rpar(leaf)

887

888

889def is_empty_lpar(leaf: Leaf) -> bool:

890 return leaf.type == token.LPAR and leaf.value == ""

891

892

893def is_empty_rpar(leaf: Leaf) -> bool:

894 return leaf.type == token.RPAR and leaf.value == ""

895

896

897def is_import(leaf: Leaf) -> bool:

898 """Return True if the given leaf starts an import statement."""

899 p = leaf.parent

900 t = leaf.type

901 v = leaf.value

902 return bool(

903 t == token.NAME

904 and (

905 (v == "import" and p and p.type == syms.import_name)

906 or (v == "from" and p and p.type == syms.import_from)

907 )

908 )

909

910

911def is_with_or_async_with_stmt(leaf: Leaf) -> bool:

912 """Return True if the given leaf starts a with or async with statement."""

913 return bool(

914 leaf.type == token.NAME

915 and leaf.value == "with"

916 and leaf.parent

917 and leaf.parent.type == syms.with_stmt

918 ) or bool(

919 leaf.type == token.ASYNC

920 and leaf.next_sibling

921 and leaf.next_sibling.type == syms.with_stmt

922 )

923

924

925def is_async_stmt_or_funcdef(leaf: Leaf) -> bool:

926 """Return True if the given leaf starts an async def/for/with statement.

927

928 Note that `async def` can be either an `async_stmt` or `async_funcdef`,

929 the latter is used when it has decorators.

930 """

931 return bool(

932 leaf.type == token.ASYNC

933 and leaf.parent

934 and leaf.parent.type in {syms.async_stmt, syms.async_funcdef}

935 )

936

937

938def is_type_comment(leaf: Leaf, mode: Mode) -> bool:

939 """Return True if the given leaf is a type comment. This function should only

940 be used for general type comments (excluding ignore annotations, which should

941 use `is_type_ignore_comment`). Note that general type comments are no longer

942 used in modern version of Python, this function may be deprecated in the future."""

943 t = leaf.type

944 v = leaf.value

945 return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_comment_string(v, mode)

946

947

948def is_type_comment_string(value: str, mode: Mode) -> bool:

949 if Preview.standardize_type_comments in mode:

950 is_valid = value.startswith("#") and value[1:].lstrip().startswith("type:")

951 else:

952 is_valid = value.startswith("# type:")

953 return is_valid

954

955

956def is_type_ignore_comment(leaf: Leaf, mode: Mode) -> bool:

957 """Return True if the given leaf is a type comment with ignore annotation."""

958 t = leaf.type

959 v = leaf.value

960 return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_ignore_comment_string(

961 v, mode

962 )

963

964

965def is_type_ignore_comment_string(value: str, mode: Mode) -> bool:

966 """Return True if the given string match with type comment with

967 ignore annotation."""

968 if Preview.standardize_type_comments in mode:

969 is_valid = is_type_comment_string(value, mode) and value.split(":", 1)[

970 1

971 ].lstrip().startswith("ignore")

972 else:

973 is_valid = value.startswith("# type: ignore")

974

975 return is_valid

976

977

978def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:

979 """Wrap `child` in parentheses.

980

981 This replaces `child` with an atom holding the parentheses and the old

982 child. That requires moving the prefix.

983

984 If `visible` is False, the leaves will be valueless (and thus invisible).

985 """

986 lpar = Leaf(token.LPAR, "(" if visible else "")

987 rpar = Leaf(token.RPAR, ")" if visible else "")

988 prefix = child.prefix

989 child.prefix = ""

990 index = child.remove() or 0

991 new_child = Node(syms.atom, [lpar, child, rpar])

992 new_child.prefix = prefix

993 parent.insert_child(index, new_child)

994

995

996def unwrap_singleton_parenthesis(node: LN) -> Optional[LN]:

997 """Returns `wrapped` if `node` is of the shape ( wrapped ).

998

999 Parenthesis can be optional. Returns None otherwise"""

1000 if len(node.children) != 3:

1001 return None

1002

1003 lpar, wrapped, rpar = node.children

1004 if not (lpar.type == token.LPAR and rpar.type == token.RPAR):

1005 return None

1006

1007 return wrapped

1008

1009

1010def ensure_visible(leaf: Leaf) -> None:

1011 """Make sure parentheses are visible.

1012

1013 They could be invisible as part of some statements (see

1014 :func:`normalize_invisible_parens` and :func:`visit_import_from`).

1015 """

1016 if leaf.type == token.LPAR:

1017 leaf.value = "("

1018 elif leaf.type == token.RPAR:

1019 leaf.value = ")"

1020

1021

1022def is_name_token(nl: NL) -> TypeGuard[Leaf]:

1023 return nl.type == token.NAME

1024

1025

1026def is_lpar_token(nl: NL) -> TypeGuard[Leaf]:

1027 return nl.type == token.LPAR

1028

1029

1030def is_rpar_token(nl: NL) -> TypeGuard[Leaf]:

1031 return nl.type == token.RPAR

1032

1033

1034def is_number_token(nl: NL) -> TypeGuard[Leaf]:

1035 return nl.type == token.NUMBER

1036

1037

1038def get_annotation_type(leaf: Leaf) -> Literal["return", "param", None]:

1039 """Returns the type of annotation this leaf is part of, if any."""

1040 ancestor = leaf.parent

1041 while ancestor is not None:

1042 if ancestor.prev_sibling and ancestor.prev_sibling.type == token.RARROW:

1043 return "return"

1044 if ancestor.parent and ancestor.parent.type == syms.tname:

1045 return "param"

1046 ancestor = ancestor.parent

1047 return None

1048

1049

1050def is_part_of_annotation(leaf: Leaf) -> bool:

1051 """Returns whether this leaf is part of a type annotation."""

1052 assert leaf.parent is not None

1053 return get_annotation_type(leaf) is not None

1054

1055

1056def first_leaf(node: LN) -> Optional[Leaf]:

1057 """Returns the first leaf of the ancestor node."""

1058 if isinstance(node, Leaf):

1059 return node

1060 elif not node.children:

1061 return None

1062 else:

1063 return first_leaf(node.children[0])

1064

1065

1066def last_leaf(node: LN) -> Optional[Leaf]:

1067 """Returns the last leaf of the ancestor node."""

1068 if isinstance(node, Leaf):

1069 return node

1070 elif not node.children:

1071 return None

1072 else:

1073 return last_leaf(node.children[-1])

1074

1075

1076def furthest_ancestor_with_last_leaf(leaf: Leaf) -> LN:

1077 """Returns the furthest ancestor that has this leaf node as the last leaf."""

1078 node: LN = leaf

1079 while node.parent and node.parent.children and node is node.parent.children[-1]:

1080 node = node.parent

1081 return node

1082

1083

1084def has_sibling_with_type(node: LN, type: int) -> bool:

1085 # Check previous siblings

1086 sibling = node.prev_sibling

1087 while sibling is not None:

1088 if sibling.type == type:

1089 return True

1090 sibling = sibling.prev_sibling

1091

1092 # Check next siblings

1093 sibling = node.next_sibling

1094 while sibling is not None:

1095 if sibling.type == type:

1096 return True

1097 sibling = sibling.next_sibling

1098

1099 return False