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

1"""

2blib2to3 Node/Leaf transformation-related utility functions.

3"""

5from collections.abc import Iterator

6from typing import Final, Generic, Literal, TypeGuard, TypeVar, Union

8from mypy_extensions import mypyc_attr

10from black.cache import CACHE_DIR

11from black.mode import Mode, Preview

12from black.strings import get_string_prefix, has_triple_quotes

13from blib2to3 import pygram

14from blib2to3.pgen2 import token

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

17pygram.initialize(CACHE_DIR)

18syms: Final = pygram.python_symbols

21# types

22T = TypeVar("T")

23LN = Union[Leaf, Node]

24LeafID = int

25NodeType = int

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

29STATEMENT: Final = {

30 syms.if_stmt,

31 syms.while_stmt,

32 syms.for_stmt,

33 syms.try_stmt,

34 syms.except_clause,

35 syms.with_stmt,

36 syms.funcdef,

37 syms.classdef,

38 syms.match_stmt,

39 syms.case_block,

40}

41STANDALONE_COMMENT: Final = 153

42token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"

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

44COMPARATORS: Final = {

45 token.LESS,

46 token.GREATER,

47 token.EQEQUAL,

48 token.NOTEQUAL,

49 token.LESSEQUAL,

50 token.GREATEREQUAL,

51}

52MATH_OPERATORS: Final = {

53 token.VBAR,

54 token.CIRCUMFLEX,

55 token.AMPER,

56 token.LEFTSHIFT,

57 token.RIGHTSHIFT,

58 token.PLUS,

59 token.MINUS,

60 token.STAR,

61 token.SLASH,

62 token.DOUBLESLASH,

63 token.PERCENT,

64 token.AT,

65 token.TILDE,

66 token.DOUBLESTAR,

67}

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

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

70VARARGS_PARENTS: Final = {

71 syms.arglist,

72 syms.argument, # double star in arglist

73 syms.trailer, # single argument to call

74 syms.typedargslist,

75 syms.varargslist, # lambdas

76}

77UNPACKING_PARENTS: Final = {

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

79 syms.dictsetmaker,

80 syms.listmaker,

81 syms.testlist_gexp,

82 syms.testlist_star_expr,

83 syms.subject_expr,

84 syms.pattern,

85}

86TEST_DESCENDANTS: Final = {

87 syms.test,

88 syms.lambdef,

89 syms.or_test,

90 syms.and_test,

91 syms.not_test,

92 syms.comparison,

93 syms.star_expr,

94 syms.expr,

95 syms.xor_expr,

96 syms.and_expr,

97 syms.shift_expr,

98 syms.arith_expr,

99 syms.trailer,

100 syms.term,

101 syms.power,

102 syms.namedexpr_test,

103}

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

105ASSIGNMENTS: Final = {

106 "=",

107 "+=",

108 "-=",

109 "*=",

110 "@=",

111 "/=",

112 "%=",

113 "&=",

114 "|=",

115 "^=",

116 "<<=",

117 ">>=",

118 "**=",

119 "//=",

120 ":",

121}

122

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

124BRACKET: Final = {

125 token.LPAR: token.RPAR,

126 token.LSQB: token.RSQB,

127 token.LBRACE: token.RBRACE,

128}

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

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

131BRACKETS: Final = OPENING_BRACKETS | CLOSING_BRACKETS

132ALWAYS_NO_SPACE: Final = CLOSING_BRACKETS | {

133 token.COMMA,

134 STANDALONE_COMMENT,

135 token.FSTRING_MIDDLE,

136 token.FSTRING_END,

137 token.TSTRING_MIDDLE,

138 token.TSTRING_END,

139 token.BANG,

140}

141

142RARROW = 55

143

144

145@mypyc_attr(allow_interpreted_subclasses=True)

146class Visitor(Generic[T]):

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

148

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

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

151

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

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

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

155 instead.

156

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

158 """

159 if node.type < 256:

160 name = token.tok_name[node.type]

161 else:

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

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

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

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

166 # generate a native call to visit_default.

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

168 if visitf:

169 yield from visitf(node)

170 else:

171 yield from self.visit_default(node)

172

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

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

175 if isinstance(node, Node):

176 for child in node.children:

177 yield from self.visit(child)

178

179

180def whitespace(leaf: Leaf, *, complex_subscript: bool, mode: Mode) -> str:

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

182

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

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

185 """

186 NO: Final[str] = ""

187 SPACE: Final[str] = " "

188 DOUBLESPACE: Final[str] = " "

189 t = leaf.type

190 p = leaf.parent

191 v = leaf.value

192 if t in ALWAYS_NO_SPACE:

193 return NO

194

195 if t == token.COMMENT:

196 return DOUBLESPACE

197

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

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

200 syms.subscript,

201 syms.subscriptlist,

202 syms.sliceop,

203 }:

204 return NO

205

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

207 syms.fstring_replacement_field,

208 syms.tstring_replacement_field,

209 ):

210 return NO

211

212 prev = leaf.prev_sibling

213 if not prev:

214 prevp = preceding_leaf(p)

215 if not prevp or prevp.type in OPENING_BRACKETS:

216 return NO

217

218 if t == token.COLON:

219 if prevp.type == token.COLON:

220 return NO

221

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

223 return NO

224

225 return SPACE

226

227 if prevp.type == token.EQUAL:

228 if prevp.parent:

229 if prevp.parent.type in {

230 syms.arglist,

231 syms.argument,

232 syms.parameters,

233 syms.varargslist,

234 }:

235 return NO

236

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

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

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

240 # that, too.

241 return prevp.prefix

242

243 elif (

244 prevp.type == token.STAR

245 and parent_type(prevp) == syms.star_expr

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

247 ):

248 # No space between typevar tuples or unpacking them.

249 return NO

250

251 elif prevp.type in VARARGS_SPECIALS:

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

253 return NO

254

255 elif prevp.type == token.COLON:

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

257 return SPACE if complex_subscript else NO

258

259 elif (

260 prevp.parent

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

262 and prevp.type in MATH_OPERATORS

263 ):

264 return NO

265

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

267 # no space in decorators

268 return NO

269

270 elif prev.type in OPENING_BRACKETS:

271 return NO

272

273 elif prev.type == token.BANG:

274 return NO

275

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

277 # untyped function signatures or calls

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

279 return NO

280

281 elif p.type == syms.varargslist:

282 # lambdas

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

284 return NO

285

286 elif p.type == syms.typedargslist:

287 # typed function signatures

288 if not prev:

289 return NO

290

291 if t == token.EQUAL:

292 if prev.type not in TYPED_NAMES:

293 return NO

294

295 elif prev.type == token.EQUAL:

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

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

298 return prev.prefix

299

300 elif prev.type != token.COMMA:

301 return NO

302

303 elif p.type in TYPED_NAMES:

304 # type names

305 if not prev:

306 prevp = preceding_leaf(p)

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

308 return NO

309

310 elif p.type == syms.trailer:

311 # attributes and calls

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

313 return NO

314

315 if not prev:

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

317 return NO

318

319 elif prev.type != token.COMMA:

320 return NO

321

322 elif p.type == syms.argument:

323 # single argument

324 if t == token.EQUAL:

325 return NO

326

327 if not prev:

328 prevp = preceding_leaf(p)

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

330 return NO

331

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

333 return NO

334

335 elif p.type == syms.decorator:

336 # decorators

337 return NO

338

339 elif p.type == syms.dotted_name:

340 if prev:

341 return NO

342

343 prevp = preceding_leaf(p)

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

345 return NO

346

347 elif p.type == syms.classdef:

348 if t == token.LPAR:

349 return NO

350

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

352 return NO

353

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

355 # indexing

356 if not prev:

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

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

359 return SPACE

360

361 return NO

362

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

364 return SPACE

365

366 elif not complex_subscript:

367 return NO

368

369 elif p.type == syms.atom:

370 if prev and t == token.DOT:

371 # dots, but not the first one.

372 return NO

373

374 elif p.type == syms.dictsetmaker:

375 # dict unpacking

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

377 return NO

378

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

380 # unary ops

381 if not prev:

382 prevp = preceding_leaf(p)

383 if not prevp or prevp.type in OPENING_BRACKETS:

384 return NO

385

386 prevp_parent = prevp.parent

387 assert prevp_parent is not None

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

389 syms.subscript,

390 syms.sliceop,

391 }:

392 return NO

393

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

395 return NO

396

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

398 return NO

399

400 elif p.type == syms.import_from:

401 if t == token.DOT:

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

403 return NO

404

405 elif t == token.NAME:

406 if v == "import":

407 return SPACE

408

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

410 return NO

411

412 elif p.type == syms.sliceop:

413 return NO

414

415 elif p.type == syms.except_clause:

416 if t == token.STAR:

417 return NO

418

419 if Preview.simplify_power_operator_hugging in mode:

420 # Power operator hugging

421 if t == token.DOUBLESTAR and is_simple_exponentiation(p):

422 return NO

423 prevp = preceding_leaf(leaf)

424 if prevp and prevp.type == token.DOUBLESTAR:

425 if prevp.parent and is_simple_exponentiation(prevp.parent):

426 return NO

427

428 return SPACE

429

430

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

432 """Generate a normalized prefix string."""

433 if form_feed:

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

435 return empty_line * nl_count

436

437

438def preceding_leaf(node: LN | None) -> Leaf | None:

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

440 while node:

441 res = node.prev_sibling

442 if res:

443 if isinstance(res, Leaf):

444 return res

445

446 try:

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

448

449 except IndexError:

450 return None

451

452 node = node.parent

453 return None

454

455

456def prev_siblings_are(node: LN | None, tokens: list[NodeType | None]) -> bool:

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

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

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

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

461 if not tokens:

462 return True

463 if tokens[-1] is None:

464 return node is None

465 if not node:

466 return False

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

468 return False

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

470

471

472def parent_type(node: LN | None) -> NodeType | None:

473 """

474 Returns:

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

476 OR

477 None, otherwise.

478 """

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

480 return None

481

482 return node.parent.type

483

484

485def child_towards(ancestor: Node, descendant: LN) -> LN | None:

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

487 node: LN | None = descendant

488 while node and node.parent != ancestor:

489 node = node.parent

490 return node

491

492

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

494 """

495 Side Effects:

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

497 @old_child's underlying Node structure.

498 OR

499 * Otherwise, this function does nothing.

500 """

501 parent = old_child.parent

502 if not parent:

503 return

504

505 child_idx = old_child.remove()

506 if child_idx is not None:

507 parent.insert_child(child_idx, new_child)

508

509

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

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

512

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

514 """

515 same_prefix = leaf.prefix

516 container: LN = leaf

517 while container:

518 parent = container.parent

519 if parent is None:

520 break

521

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

523 break

524

525 if parent.type == syms.file_input:

526 break

527

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

529 break

530

531 container = parent

532 return container

533

534

535def first_leaf_of(node: LN) -> Leaf | None:

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

537 if isinstance(node, Leaf):

538 return node

539 if node.children:

540 return first_leaf_of(node.children[0])

541 else:

542 return None

543

544

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

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

547 return node.type in {

548 syms.arith_expr,

549 syms.shift_expr,

550 syms.xor_expr,

551 syms.and_expr,

552 }

553

554

555def is_simple_exponentiation(node: LN) -> bool:

556 """Whether whitespace around `**` should be removed."""

557

558 def is_simple(node: LN) -> bool:

559 if isinstance(node, Leaf):

560 return node.type in (token.NAME, token.NUMBER, token.DOT, token.DOUBLESTAR)

561 elif node.type == syms.factor: # unary operators

562 return is_simple(node.children[1])

563 else:

564 return all(is_simple(child) for child in node.children)

565

566 return (

567 node.type == syms.power

568 and len(node.children) >= 3

569 and node.children[-2].type == token.DOUBLESTAR

570 and is_simple(node)

571 )

572

573

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

575 if isinstance(node, Leaf):

576 if node.type != token.STRING:

577 return False

578

579 prefix = get_string_prefix(node.value)

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

581 return False

582

583 if (

584 node.parent

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

586 and not node.parent.prev_sibling

587 and node.parent.parent

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

589 ):

590 return True

591

592 if prev_siblings_are(

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

594 ):

595 return True

596

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

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

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

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

601 return True

602

603 return False

604

605

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

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

608 return (

609 node.type == syms.atom

610 and len(node.children) == 2

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

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

613 )

614

615

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

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

618 if node.type == syms.atom:

619 gexp = unwrap_singleton_parenthesis(node)

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

621 return False

622

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

624

625 return (

626 node.type in IMPLICIT_TUPLE

627 and len(node.children) == 2

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

629 )

630

631

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

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

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 True

641

642

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

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

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.namedexpr_test for child in gexp.children)

652

653

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

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

656 if node.type != syms.atom:

657 return False

658 gexp = unwrap_singleton_parenthesis(node)

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

660 return False

661

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

663

664

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

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

667 if node.type != syms.atom:

668 return False

669 gexp = unwrap_singleton_parenthesis(node)

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

671 return False

672

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

674

675

676def is_one_sequence_between(

677 opening: Leaf,

678 closing: Leaf,

679 leaves: list[Leaf],

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

681) -> bool:

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

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

684 return False

685

686 depth = closing.bracket_depth + 1

687 for _opening_index, leaf in enumerate(leaves):

688 if leaf is opening:

689 break

690

691 else:

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

693

694 commas = 0

695 _opening_index += 1

696 for leaf in leaves[_opening_index:]:

697 if leaf is closing:

698 break

699

700 bracket_depth = leaf.bracket_depth

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

702 commas += 1

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

704 syms.arglist,

705 syms.typedargslist,

706 }:

707 commas += 1

708 break

709

710 return commas < 2

711

712

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

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

715 inner = unwrap_singleton_parenthesis(node)

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

717

718

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

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

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

722 (

723 len(node.children) == 2

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

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

726 )

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

728 or (

729 last

730 and len(node.children) == 2

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

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

733 )

734 # last trailer can be arguments

735 or (

736 last

737 and len(node.children) == 3

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

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

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

741 )

742 )

743

744

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

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

747

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

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

750

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

752 The new grammar is : decorator: @ namedexpr_test NEWLINE

753 """

754 if node.type == token.NAME:

755 return True

756 if node.type == syms.power:

757 if node.children:

758 return (

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

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

761 and (

762 len(node.children) < 2

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

764 )

765 )

766 return False

767

768

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

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

771 if node.type == syms.yield_expr:

772 return True

773

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

775 return True

776

777 if node.type != syms.atom:

778 return False

779

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

781 return False

782

783 lpar, expr, rpar = node.children

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

785 return is_yield(expr)

786

787 return False

788

789

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

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

792

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

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

795 extended iterable unpacking (PEP 3132) and additional unpacking

796 generalizations (PEP 448).

797 """

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

799 return False

800

801 p = leaf.parent

802 if p.type == syms.star_expr:

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

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

805 if not p.parent:

806 return False

807

808 p = p.parent

809

810 return p.type in within

811

812

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

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

815 return node.type == syms.fstring

816

817

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

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

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

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

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

823 return string_leaf

824

825

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

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

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

829 leaf = fstring_tstring_to_string(node)

830 elif isinstance(node, Leaf):

831 leaf = node

832 else:

833 return False

834

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

836

837

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

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

840 assert node.parent is not None

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

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

843

844

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

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

847

848

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

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

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

852 return False

853

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

855 if node.prefix.strip():

856 return False

857

858 if (

859 len(node.children) != 4

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

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

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

863 ):

864 return False

865

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

867 return False

868

869 return is_stub_body(node.children[2])

870

871

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

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

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

875 return False

876

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

878 return False

879

880 child = node.children[0]

881 return (

882 not child.prefix.strip()

883 and child.type == syms.atom

884 and len(child.children) == 3

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

886 )

887

888

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

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

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

892 """

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

894 return False

895

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

897 return (

898 isinstance(first, Leaf)

899 and first.type == token.LPAR

900 and first.value == ""

901 and isinstance(last, Leaf)

902 and last.type == token.RPAR

903 and last.value == ""

904 )

905

906

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

908 return is_empty_lpar(leaf) or is_empty_rpar(leaf)

909

910

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

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

913

914

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

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

917

918

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

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

921 p = leaf.parent

922 t = leaf.type

923 v = leaf.value

924 return bool(

925 t == token.NAME

926 and (

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

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

929 )

930 )

931

932

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

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

935 return bool(

936 leaf.type == token.NAME

937 and leaf.value == "with"

938 and leaf.parent

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

940 ) or bool(

941 leaf.type == token.ASYNC

942 and leaf.next_sibling

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

944 )

945

946

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

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

949

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

951 the latter is used when it has decorators.

952 """

953 return bool(

954 leaf.type == token.ASYNC

955 and leaf.parent

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

957 )

958

959

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

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

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

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

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

965 t = leaf.type

966 v = leaf.value

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

968

969

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

971 return value.startswith("#") and value[1:].lstrip().startswith("type:")

972

973

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

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

976 t = leaf.type

977 v = leaf.value

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

979 v, mode

980 )

981

982

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

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

985 ignore annotation."""

986 return is_type_comment_string(value, mode) and value.split(":", 1)[

987 1

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

989

990

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

992 """Wrap `child` in parentheses.

993

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

995 child. That requires moving the prefix.

996

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

998 """

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

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

1001 prefix = child.prefix

1002 child.prefix = ""

1003 index = child.remove() or 0

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

1005 new_child.prefix = prefix

1006 parent.insert_child(index, new_child)

1007

1008

1009def unwrap_singleton_parenthesis(node: LN) -> LN | None:

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

1011

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

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

1014 return None

1015

1016 lpar, wrapped, rpar = node.children

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

1018 return None

1019

1020 return wrapped

1021

1022

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

1024 """Make sure parentheses are visible.

1025

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

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

1028 """

1029 if leaf.type == token.LPAR:

1030 leaf.value = "("

1031 elif leaf.type == token.RPAR:

1032 leaf.value = ")"

1033

1034

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

1036 return nl.type == token.NAME

1037

1038

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

1040 return nl.type == token.LPAR

1041

1042

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

1044 return nl.type == token.RPAR

1045

1046

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

1048 return nl.type == token.NUMBER

1049

1050

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

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

1053 ancestor = leaf.parent

1054 while ancestor is not None:

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

1056 return "return"

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

1058 return "param"

1059 ancestor = ancestor.parent

1060 return None

1061

1062

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

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

1065 assert leaf.parent is not None

1066 return get_annotation_type(leaf) is not None

1067

1068

1069def first_leaf(node: LN) -> Leaf | None:

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

1071 if isinstance(node, Leaf):

1072 return node

1073 elif not node.children:

1074 return None

1075 else:

1076 return first_leaf(node.children[0])

1077

1078

1079def last_leaf(node: LN) -> Leaf | None:

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

1081 if isinstance(node, Leaf):

1082 return node

1083 elif not node.children:

1084 return None

1085 else:

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

1087

1088

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

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

1091 node: LN = leaf

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

1093 node = node.parent

1094 return node

1095

1096

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

1098 # Check previous siblings

1099 sibling = node.prev_sibling

1100 while sibling is not None:

1101 if sibling.type == type:

1102 return True

1103 sibling = sibling.prev_sibling

1104

1105 # Check next siblings

1106 sibling = node.next_sibling

1107 while sibling is not None:

1108 if sibling.type == type:

1109 return True

1110 sibling = sibling.next_sibling

1111

1112 return False