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

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.BANG,

144}

145

146RARROW = 55

147

148

149@mypyc_attr(allow_interpreted_subclasses=True)

150class Visitor(Generic[T]):

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

152

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

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

155

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

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

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

159 instead.

160

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

162 """

163 if node.type < 256:

164 name = token.tok_name[node.type]

165 else:

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

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

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

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

170 # generate a native call to visit_default.

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

172 if visitf:

173 yield from visitf(node)

174 else:

175 yield from self.visit_default(node)

176

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

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

179 if isinstance(node, Node):

180 for child in node.children:

181 yield from self.visit(child)

182

183

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

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

186

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

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

189 """

190 NO: Final[str] = ""

191 SPACE: Final[str] = " "

192 DOUBLESPACE: Final[str] = " "

193 t = leaf.type

194 p = leaf.parent

195 v = leaf.value

196 if t in ALWAYS_NO_SPACE:

197 return NO

198

199 if t == token.COMMENT:

200 return DOUBLESPACE

201

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

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

204 syms.subscript,

205 syms.subscriptlist,

206 syms.sliceop,

207 }:

208 return NO

209

210 if t == token.LBRACE and p.type == syms.fstring_replacement_field:

211 return NO

212

213 prev = leaf.prev_sibling

214 if not prev:

215 prevp = preceding_leaf(p)

216 if not prevp or prevp.type in OPENING_BRACKETS:

217 return NO

218

219 if t == token.COLON:

220 if prevp.type == token.COLON:

221 return NO

222

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

224 return NO

225

226 return SPACE

227

228 if prevp.type == token.EQUAL:

229 if prevp.parent:

230 if prevp.parent.type in {

231 syms.arglist,

232 syms.argument,

233 syms.parameters,

234 syms.varargslist,

235 }:

236 return NO

237

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

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

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

241 # that, too.

242 return prevp.prefix

243

244 elif (

245 prevp.type == token.STAR

246 and parent_type(prevp) == syms.star_expr

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

248 ):

249 # No space between typevar tuples or unpacking them.

250 return NO

251

252 elif prevp.type in VARARGS_SPECIALS:

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

254 return NO

255

256 elif prevp.type == token.COLON:

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

258 return SPACE if complex_subscript else NO

259

260 elif (

261 prevp.parent

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

263 and prevp.type in MATH_OPERATORS

264 ):

265 return NO

266

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

268 # no space in decorators

269 return NO

270

271 elif prev.type in OPENING_BRACKETS:

272 return NO

273

274 elif prev.type == token.BANG:

275 return NO

276

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

278 # untyped function signatures or calls

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

280 return NO

281

282 elif p.type == syms.varargslist:

283 # lambdas

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

285 return NO

286

287 elif p.type == syms.typedargslist:

288 # typed function signatures

289 if not prev:

290 return NO

291

292 if t == token.EQUAL:

293 if prev.type not in TYPED_NAMES:

294 return NO

295

296 elif prev.type == token.EQUAL:

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

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

299 return prev.prefix

300

301 elif prev.type != token.COMMA:

302 return NO

303

304 elif p.type in TYPED_NAMES:

305 # type names

306 if not prev:

307 prevp = preceding_leaf(p)

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

309 return NO

310

311 elif p.type == syms.trailer:

312 # attributes and calls

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

314 return NO

315

316 if not prev:

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

318 return NO

319

320 elif prev.type != token.COMMA:

321 return NO

322

323 elif p.type == syms.argument:

324 # single argument

325 if t == token.EQUAL:

326 return NO

327

328 if not prev:

329 prevp = preceding_leaf(p)

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

331 return NO

332

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

334 return NO

335

336 elif p.type == syms.decorator:

337 # decorators

338 return NO

339

340 elif p.type == syms.dotted_name:

341 if prev:

342 return NO

343

344 prevp = preceding_leaf(p)

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

346 return NO

347

348 elif p.type == syms.classdef:

349 if t == token.LPAR:

350 return NO

351

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

353 return NO

354

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

356 # indexing

357 if not prev:

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

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

360 return SPACE

361

362 return NO

363

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

365 return SPACE

366

367 elif not complex_subscript:

368 return NO

369

370 elif p.type == syms.atom:

371 if prev and t == token.DOT:

372 # dots, but not the first one.

373 return NO

374

375 elif p.type == syms.dictsetmaker:

376 # dict unpacking

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

378 return NO

379

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

381 # unary ops

382 if not prev:

383 prevp = preceding_leaf(p)

384 if not prevp or prevp.type in OPENING_BRACKETS:

385 return NO

386

387 prevp_parent = prevp.parent

388 assert prevp_parent is not None

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

390 syms.subscript,

391 syms.sliceop,

392 }:

393 return NO

394

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

396 return NO

397

398 # TODO: add fstring here?

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

400 return NO

401

402 elif p.type == syms.import_from:

403 if t == token.DOT:

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

405 return NO

406

407 elif t == token.NAME:

408 if v == "import":

409 return SPACE

410

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

412 return NO

413

414 elif p.type == syms.sliceop:

415 return NO

416

417 elif p.type == syms.except_clause:

418 if t == token.STAR:

419 return NO

420

421 return SPACE

422

423

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

425 """Generate a normalized prefix string."""

426 if form_feed:

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

428 return empty_line * nl_count

429

430

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

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

433 while node:

434 res = node.prev_sibling

435 if res:

436 if isinstance(res, Leaf):

437 return res

438

439 try:

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

441

442 except IndexError:

443 return None

444

445 node = node.parent

446 return None

447

448

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

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

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

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

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

454 if not tokens:

455 return True

456 if tokens[-1] is None:

457 return node is None

458 if not node:

459 return False

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

461 return False

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

463

464

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

466 """

467 Returns:

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

469 OR

470 None, otherwise.

471 """

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

473 return None

474

475 return node.parent.type

476

477

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

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

480 node: Optional[LN] = descendant

481 while node and node.parent != ancestor:

482 node = node.parent

483 return node

484

485

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

487 """

488 Side Effects:

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

490 @old_child's underlying Node structure.

491 OR

492 * Otherwise, this function does nothing.

493 """

494 parent = old_child.parent

495 if not parent:

496 return

497

498 child_idx = old_child.remove()

499 if child_idx is not None:

500 parent.insert_child(child_idx, new_child)

501

502

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

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

505

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

507 """

508 same_prefix = leaf.prefix

509 container: LN = leaf

510 while container:

511 parent = container.parent

512 if parent is None:

513 break

514

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

516 break

517

518 if parent.type == syms.file_input:

519 break

520

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

522 break

523

524 container = parent

525 return container

526

527

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

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

530 if isinstance(node, Leaf):

531 return node

532 if node.children:

533 return first_leaf_of(node.children[0])

534 else:

535 return None

536

537

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

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

540 return node.type in {

541 syms.arith_expr,

542 syms.shift_expr,

543 syms.xor_expr,

544 syms.and_expr,

545 }

546

547

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

549 if isinstance(node, Leaf):

550 if node.type != token.STRING:

551 return False

552

553 prefix = get_string_prefix(node.value)

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

555 return False

556

557 if (

558 node.parent

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

560 and not node.parent.prev_sibling

561 and node.parent.parent

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

563 ):

564 return True

565

566 if prev_siblings_are(

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

568 ):

569 return True

570

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

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

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

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

575 return True

576

577 return False

578

579

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

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

582 return (

583 node.type == syms.atom

584 and len(node.children) == 2

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

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

587 )

588

589

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

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

592 if node.type == syms.atom:

593 gexp = unwrap_singleton_parenthesis(node)

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

595 return False

596

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

598

599 return (

600 node.type in IMPLICIT_TUPLE

601 and len(node.children) == 2

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

603 )

604

605

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

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

608 if node.type != syms.atom:

609 return False

610 gexp = unwrap_singleton_parenthesis(node)

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

612 return False

613

614 return True

615

616

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

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

619 if node.type != syms.atom:

620 return False

621 gexp = unwrap_singleton_parenthesis(node)

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

623 return False

624

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

626

627

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

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

630 if node.type != syms.atom:

631 return False

632 gexp = unwrap_singleton_parenthesis(node)

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

634 return False

635

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

637

638

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

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

641 if node.type != syms.atom:

642 return False

643 gexp = unwrap_singleton_parenthesis(node)

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

645 return False

646

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

648

649

650def is_one_sequence_between(

651 opening: Leaf,

652 closing: Leaf,

653 leaves: list[Leaf],

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

655) -> bool:

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

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

658 return False

659

660 depth = closing.bracket_depth + 1

661 for _opening_index, leaf in enumerate(leaves):

662 if leaf is opening:

663 break

664

665 else:

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

667

668 commas = 0

669 _opening_index += 1

670 for leaf in leaves[_opening_index:]:

671 if leaf is closing:

672 break

673

674 bracket_depth = leaf.bracket_depth

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

676 commas += 1

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

678 syms.arglist,

679 syms.typedargslist,

680 }:

681 commas += 1

682 break

683

684 return commas < 2

685

686

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

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

689 inner = unwrap_singleton_parenthesis(node)

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

691

692

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

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

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

696 (

697 len(node.children) == 2

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

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

700 )

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

702 or (

703 last

704 and len(node.children) == 2

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

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

707 )

708 # last trailer can be arguments

709 or (

710 last

711 and len(node.children) == 3

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

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

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

715 )

716 )

717

718

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

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

721

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

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

724

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

726 The new grammar is : decorator: @ namedexpr_test NEWLINE

727 """

728 if node.type == token.NAME:

729 return True

730 if node.type == syms.power:

731 if node.children:

732 return (

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

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

735 and (

736 len(node.children) < 2

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

738 )

739 )

740 return False

741

742

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

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

745 if node.type == syms.yield_expr:

746 return True

747

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

749 return True

750

751 if node.type != syms.atom:

752 return False

753

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

755 return False

756

757 lpar, expr, rpar = node.children

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

759 return is_yield(expr)

760

761 return False

762

763

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

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

766

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

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

769 extended iterable unpacking (PEP 3132) and additional unpacking

770 generalizations (PEP 448).

771 """

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

773 return False

774

775 p = leaf.parent

776 if p.type == syms.star_expr:

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

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

779 if not p.parent:

780 return False

781

782 p = p.parent

783

784 return p.type in within

785

786

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

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

789 return node.type == syms.fstring

790

791

792def fstring_to_string(node: Node) -> Leaf:

793 """Converts an fstring node back to a string node."""

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

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

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

797 return string_leaf

798

799

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

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

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

803 leaf = fstring_to_string(node)

804 elif isinstance(node, Leaf):

805 leaf = node

806 else:

807 return False

808

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

810

811

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

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

814 assert node.parent is not None

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

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

817

818

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

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

821

822

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

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

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

826 return False

827

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

829 if node.prefix.strip():

830 return False

831

832 if (

833 len(node.children) != 4

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

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

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

837 ):

838 return False

839

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

841 return False

842

843 return is_stub_body(node.children[2])

844

845

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

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

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

849 return False

850

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

852 return False

853

854 child = node.children[0]

855 return (

856 not child.prefix.strip()

857 and child.type == syms.atom

858 and len(child.children) == 3

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

860 )

861

862

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

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

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

866 """

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

868 return False

869

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

871 return (

872 isinstance(first, Leaf)

873 and first.type == token.LPAR

874 and first.value == ""

875 and isinstance(last, Leaf)

876 and last.type == token.RPAR

877 and last.value == ""

878 )

879

880

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

882 return is_empty_lpar(leaf) or is_empty_rpar(leaf)

883

884

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

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

887

888

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

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

891

892

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

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

895 p = leaf.parent

896 t = leaf.type

897 v = leaf.value

898 return bool(

899 t == token.NAME

900 and (

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

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

903 )

904 )

905

906

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

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

909 return bool(

910 leaf.type == token.NAME

911 and leaf.value == "with"

912 and leaf.parent

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

914 ) or bool(

915 leaf.type == token.ASYNC

916 and leaf.next_sibling

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

918 )

919

920

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

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

923

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

925 the latter is used when it has decorators.

926 """

927 return bool(

928 leaf.type == token.ASYNC

929 and leaf.parent

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

931 )

932

933

934def is_type_comment(leaf: Leaf) -> bool:

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

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

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

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

939 t = leaf.type

940 v = leaf.value

941 return t in {token.COMMENT, STANDALONE_COMMENT} and v.startswith("# type:")

942

943

944def is_type_ignore_comment(leaf: Leaf) -> bool:

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

946 t = leaf.type

947 v = leaf.value

948 return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_ignore_comment_string(v)

949

950

951def is_type_ignore_comment_string(value: str) -> bool:

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

953 ignore annotation."""

954 return value.startswith("# type: ignore")

955

956

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

958 """Wrap `child` in parentheses.

959

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

961 child. That requires moving the prefix.

962

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

964 """

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

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

967 prefix = child.prefix

968 child.prefix = ""

969 index = child.remove() or 0

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

971 new_child.prefix = prefix

972 parent.insert_child(index, new_child)

973

974

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

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

977

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

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

980 return None

981

982 lpar, wrapped, rpar = node.children

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

984 return None

985

986 return wrapped

987

988

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

990 """Make sure parentheses are visible.

991

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

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

994 """

995 if leaf.type == token.LPAR:

996 leaf.value = "("

997 elif leaf.type == token.RPAR:

998 leaf.value = ")"

999

1000

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

1002 return nl.type == token.NAME

1003

1004

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

1006 return nl.type == token.LPAR

1007

1008

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

1010 return nl.type == token.RPAR

1011

1012

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

1014 return nl.type == token.NUMBER

1015

1016

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

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

1019 ancestor = leaf.parent

1020 while ancestor is not None:

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

1022 return "return"

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

1024 return "param"

1025 ancestor = ancestor.parent

1026 return None

1027

1028

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

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

1031 assert leaf.parent is not None

1032 return get_annotation_type(leaf) is not None

1033

1034

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

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

1037 if isinstance(node, Leaf):

1038 return node

1039 elif not node.children:

1040 return None

1041 else:

1042 return first_leaf(node.children[0])

1043

1044

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

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

1047 if isinstance(node, Leaf):

1048 return node

1049 elif not node.children:

1050 return None

1051 else:

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

1053

1054

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

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

1057 node: LN = leaf

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

1059 node = node.parent

1060 return node

1061

1062

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

1064 # Check previous siblings

1065 sibling = node.prev_sibling

1066 while sibling is not None:

1067 if sibling.type == type:

1068 return True

1069 sibling = sibling.prev_sibling

1070

1071 # Check next siblings

1072 sibling = node.next_sibling

1073 while sibling is not None:

1074 if sibling.type == type:

1075 return True

1076 sibling = sibling.next_sibling

1077

1078 return False