Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# mako/_ast_util.py 

2# Copyright 2006-2020 the Mako authors and contributors <see AUTHORS file> 

3# 

4# This module is part of Mako and is released under 

5# the MIT License: http://www.opensource.org/licenses/mit-license.php 

6 

7""" 

8 ast 

9 ~~~ 

10 

11 This is a stripped down version of Armin Ronacher's ast module. 

12 

13 :copyright: Copyright 2008 by Armin Ronacher. 

14 :license: Python License. 

15""" 

16 

17 

18from _ast import Add 

19from _ast import And 

20from _ast import AST 

21from _ast import BitAnd 

22from _ast import BitOr 

23from _ast import BitXor 

24from _ast import Div 

25from _ast import Eq 

26from _ast import FloorDiv 

27from _ast import Gt 

28from _ast import GtE 

29from _ast import If 

30from _ast import In 

31from _ast import Invert 

32from _ast import Is 

33from _ast import IsNot 

34from _ast import LShift 

35from _ast import Lt 

36from _ast import LtE 

37from _ast import Mod 

38from _ast import Mult 

39from _ast import Name 

40from _ast import Not 

41from _ast import NotEq 

42from _ast import NotIn 

43from _ast import Or 

44from _ast import PyCF_ONLY_AST 

45from _ast import RShift 

46from _ast import Sub 

47from _ast import UAdd 

48from _ast import USub 

49 

50from mako.compat import arg_stringname 

51 

52BOOLOP_SYMBOLS = {And: "and", Or: "or"} 

53 

54BINOP_SYMBOLS = { 

55 Add: "+", 

56 Sub: "-", 

57 Mult: "*", 

58 Div: "/", 

59 FloorDiv: "//", 

60 Mod: "%", 

61 LShift: "<<", 

62 RShift: ">>", 

63 BitOr: "|", 

64 BitAnd: "&", 

65 BitXor: "^", 

66} 

67 

68CMPOP_SYMBOLS = { 

69 Eq: "==", 

70 Gt: ">", 

71 GtE: ">=", 

72 In: "in", 

73 Is: "is", 

74 IsNot: "is not", 

75 Lt: "<", 

76 LtE: "<=", 

77 NotEq: "!=", 

78 NotIn: "not in", 

79} 

80 

81UNARYOP_SYMBOLS = {Invert: "~", Not: "not", UAdd: "+", USub: "-"} 

82 

83ALL_SYMBOLS = {} 

84ALL_SYMBOLS.update(BOOLOP_SYMBOLS) 

85ALL_SYMBOLS.update(BINOP_SYMBOLS) 

86ALL_SYMBOLS.update(CMPOP_SYMBOLS) 

87ALL_SYMBOLS.update(UNARYOP_SYMBOLS) 

88 

89 

90def parse(expr, filename="<unknown>", mode="exec"): 

91 """Parse an expression into an AST node.""" 

92 return compile(expr, filename, mode, PyCF_ONLY_AST) 

93 

94 

95def iter_fields(node): 

96 """Iterate over all fields of a node, only yielding existing fields.""" 

97 # CPython 2.5 compat 

98 if not hasattr(node, "_fields") or not node._fields: 

99 return 

100 for field in node._fields: 

101 try: 

102 yield field, getattr(node, field) 

103 except AttributeError: 

104 pass 

105 

106 

107class NodeVisitor(object): 

108 

109 """ 

110 Walks the abstract syntax tree and call visitor functions for every node 

111 found. The visitor functions may return values which will be forwarded 

112 by the `visit` method. 

113 

114 Per default the visitor functions for the nodes are ``'visit_'`` + 

115 class name of the node. So a `TryFinally` node visit function would 

116 be `visit_TryFinally`. This behavior can be changed by overriding 

117 the `get_visitor` function. If no visitor function exists for a node 

118 (return value `None`) the `generic_visit` visitor is used instead. 

119 

120 Don't use the `NodeVisitor` if you want to apply changes to nodes during 

121 traversing. For this a special visitor exists (`NodeTransformer`) that 

122 allows modifications. 

123 """ 

124 

125 def get_visitor(self, node): 

126 """ 

127 Return the visitor function for this node or `None` if no visitor 

128 exists for this node. In that case the generic visit function is 

129 used instead. 

130 """ 

131 method = "visit_" + node.__class__.__name__ 

132 return getattr(self, method, None) 

133 

134 def visit(self, node): 

135 """Visit a node.""" 

136 f = self.get_visitor(node) 

137 if f is not None: 

138 return f(node) 

139 return self.generic_visit(node) 

140 

141 def generic_visit(self, node): 

142 """Called if no explicit visitor function exists for a node.""" 

143 for field, value in iter_fields(node): 

144 if isinstance(value, list): 

145 for item in value: 

146 if isinstance(item, AST): 

147 self.visit(item) 

148 elif isinstance(value, AST): 

149 self.visit(value) 

150 

151 

152class NodeTransformer(NodeVisitor): 

153 

154 """ 

155 Walks the abstract syntax tree and allows modifications of nodes. 

156 

157 The `NodeTransformer` will walk the AST and use the return value of the 

158 visitor functions to replace or remove the old node. If the return 

159 value of the visitor function is `None` the node will be removed 

160 from the previous location otherwise it's replaced with the return 

161 value. The return value may be the original node in which case no 

162 replacement takes place. 

163 

164 Here an example transformer that rewrites all `foo` to `data['foo']`:: 

165 

166 class RewriteName(NodeTransformer): 

167 

168 def visit_Name(self, node): 

169 return copy_location(Subscript( 

170 value=Name(id='data', ctx=Load()), 

171 slice=Index(value=Str(s=node.id)), 

172 ctx=node.ctx 

173 ), node) 

174 

175 Keep in mind that if the node you're operating on has child nodes 

176 you must either transform the child nodes yourself or call the generic 

177 visit function for the node first. 

178 

179 Nodes that were part of a collection of statements (that applies to 

180 all statement nodes) may also return a list of nodes rather than just 

181 a single node. 

182 

183 Usually you use the transformer like this:: 

184 

185 node = YourTransformer().visit(node) 

186 """ 

187 

188 def generic_visit(self, node): 

189 for field, old_value in iter_fields(node): 

190 old_value = getattr(node, field, None) 

191 if isinstance(old_value, list): 

192 new_values = [] 

193 for value in old_value: 

194 if isinstance(value, AST): 

195 value = self.visit(value) 

196 if value is None: 

197 continue 

198 elif not isinstance(value, AST): 

199 new_values.extend(value) 

200 continue 

201 new_values.append(value) 

202 old_value[:] = new_values 

203 elif isinstance(old_value, AST): 

204 new_node = self.visit(old_value) 

205 if new_node is None: 

206 delattr(node, field) 

207 else: 

208 setattr(node, field, new_node) 

209 return node 

210 

211 

212class SourceGenerator(NodeVisitor): 

213 

214 """ 

215 This visitor is able to transform a well formed syntax tree into python 

216 sourcecode. For more details have a look at the docstring of the 

217 `node_to_source` function. 

218 """ 

219 

220 def __init__(self, indent_with): 

221 self.result = [] 

222 self.indent_with = indent_with 

223 self.indentation = 0 

224 self.new_lines = 0 

225 

226 def write(self, x): 

227 if self.new_lines: 

228 if self.result: 

229 self.result.append("\n" * self.new_lines) 

230 self.result.append(self.indent_with * self.indentation) 

231 self.new_lines = 0 

232 self.result.append(x) 

233 

234 def newline(self, n=1): 

235 self.new_lines = max(self.new_lines, n) 

236 

237 def body(self, statements): 

238 self.new_line = True 

239 self.indentation += 1 

240 for stmt in statements: 

241 self.visit(stmt) 

242 self.indentation -= 1 

243 

244 def body_or_else(self, node): 

245 self.body(node.body) 

246 if node.orelse: 

247 self.newline() 

248 self.write("else:") 

249 self.body(node.orelse) 

250 

251 def signature(self, node): 

252 want_comma = [] 

253 

254 def write_comma(): 

255 if want_comma: 

256 self.write(", ") 

257 else: 

258 want_comma.append(True) 

259 

260 padding = [None] * (len(node.args) - len(node.defaults)) 

261 for arg, default in zip(node.args, padding + node.defaults): 

262 write_comma() 

263 self.visit(arg) 

264 if default is not None: 

265 self.write("=") 

266 self.visit(default) 

267 if node.vararg is not None: 

268 write_comma() 

269 self.write("*" + arg_stringname(node.vararg)) 

270 if node.kwarg is not None: 

271 write_comma() 

272 self.write("**" + arg_stringname(node.kwarg)) 

273 

274 def decorators(self, node): 

275 for decorator in node.decorator_list: 

276 self.newline() 

277 self.write("@") 

278 self.visit(decorator) 

279 

280 # Statements 

281 

282 def visit_Assign(self, node): 

283 self.newline() 

284 for idx, target in enumerate(node.targets): 

285 if idx: 

286 self.write(", ") 

287 self.visit(target) 

288 self.write(" = ") 

289 self.visit(node.value) 

290 

291 def visit_AugAssign(self, node): 

292 self.newline() 

293 self.visit(node.target) 

294 self.write(BINOP_SYMBOLS[type(node.op)] + "=") 

295 self.visit(node.value) 

296 

297 def visit_ImportFrom(self, node): 

298 self.newline() 

299 self.write("from %s%s import " % ("." * node.level, node.module)) 

300 for idx, item in enumerate(node.names): 

301 if idx: 

302 self.write(", ") 

303 self.write(item) 

304 

305 def visit_Import(self, node): 

306 self.newline() 

307 for item in node.names: 

308 self.write("import ") 

309 self.visit(item) 

310 

311 def visit_Expr(self, node): 

312 self.newline() 

313 self.generic_visit(node) 

314 

315 def visit_FunctionDef(self, node): 

316 self.newline(n=2) 

317 self.decorators(node) 

318 self.newline() 

319 self.write("def %s(" % node.name) 

320 self.signature(node.args) 

321 self.write("):") 

322 self.body(node.body) 

323 

324 def visit_ClassDef(self, node): 

325 have_args = [] 

326 

327 def paren_or_comma(): 

328 if have_args: 

329 self.write(", ") 

330 else: 

331 have_args.append(True) 

332 self.write("(") 

333 

334 self.newline(n=3) 

335 self.decorators(node) 

336 self.newline() 

337 self.write("class %s" % node.name) 

338 for base in node.bases: 

339 paren_or_comma() 

340 self.visit(base) 

341 # XXX: the if here is used to keep this module compatible 

342 # with python 2.6. 

343 if hasattr(node, "keywords"): 

344 for keyword in node.keywords: 

345 paren_or_comma() 

346 self.write(keyword.arg + "=") 

347 self.visit(keyword.value) 

348 if getattr(node, "starargs", None): 

349 paren_or_comma() 

350 self.write("*") 

351 self.visit(node.starargs) 

352 if getattr(node, "kwargs", None): 

353 paren_or_comma() 

354 self.write("**") 

355 self.visit(node.kwargs) 

356 self.write(have_args and "):" or ":") 

357 self.body(node.body) 

358 

359 def visit_If(self, node): 

360 self.newline() 

361 self.write("if ") 

362 self.visit(node.test) 

363 self.write(":") 

364 self.body(node.body) 

365 while True: 

366 else_ = node.orelse 

367 if len(else_) == 1 and isinstance(else_[0], If): 

368 node = else_[0] 

369 self.newline() 

370 self.write("elif ") 

371 self.visit(node.test) 

372 self.write(":") 

373 self.body(node.body) 

374 else: 

375 self.newline() 

376 self.write("else:") 

377 self.body(else_) 

378 break 

379 

380 def visit_For(self, node): 

381 self.newline() 

382 self.write("for ") 

383 self.visit(node.target) 

384 self.write(" in ") 

385 self.visit(node.iter) 

386 self.write(":") 

387 self.body_or_else(node) 

388 

389 def visit_While(self, node): 

390 self.newline() 

391 self.write("while ") 

392 self.visit(node.test) 

393 self.write(":") 

394 self.body_or_else(node) 

395 

396 def visit_With(self, node): 

397 self.newline() 

398 self.write("with ") 

399 self.visit(node.context_expr) 

400 if node.optional_vars is not None: 

401 self.write(" as ") 

402 self.visit(node.optional_vars) 

403 self.write(":") 

404 self.body(node.body) 

405 

406 def visit_Pass(self, node): 

407 self.newline() 

408 self.write("pass") 

409 

410 def visit_Print(self, node): 

411 # XXX: python 2.6 only 

412 self.newline() 

413 self.write("print ") 

414 want_comma = False 

415 if node.dest is not None: 

416 self.write(" >> ") 

417 self.visit(node.dest) 

418 want_comma = True 

419 for value in node.values: 

420 if want_comma: 

421 self.write(", ") 

422 self.visit(value) 

423 want_comma = True 

424 if not node.nl: 

425 self.write(",") 

426 

427 def visit_Delete(self, node): 

428 self.newline() 

429 self.write("del ") 

430 for idx, target in enumerate(node): 

431 if idx: 

432 self.write(", ") 

433 self.visit(target) 

434 

435 def visit_TryExcept(self, node): 

436 self.newline() 

437 self.write("try:") 

438 self.body(node.body) 

439 for handler in node.handlers: 

440 self.visit(handler) 

441 

442 def visit_TryFinally(self, node): 

443 self.newline() 

444 self.write("try:") 

445 self.body(node.body) 

446 self.newline() 

447 self.write("finally:") 

448 self.body(node.finalbody) 

449 

450 def visit_Global(self, node): 

451 self.newline() 

452 self.write("global " + ", ".join(node.names)) 

453 

454 def visit_Nonlocal(self, node): 

455 self.newline() 

456 self.write("nonlocal " + ", ".join(node.names)) 

457 

458 def visit_Return(self, node): 

459 self.newline() 

460 self.write("return ") 

461 self.visit(node.value) 

462 

463 def visit_Break(self, node): 

464 self.newline() 

465 self.write("break") 

466 

467 def visit_Continue(self, node): 

468 self.newline() 

469 self.write("continue") 

470 

471 def visit_Raise(self, node): 

472 # XXX: Python 2.6 / 3.0 compatibility 

473 self.newline() 

474 self.write("raise") 

475 if hasattr(node, "exc") and node.exc is not None: 

476 self.write(" ") 

477 self.visit(node.exc) 

478 if node.cause is not None: 

479 self.write(" from ") 

480 self.visit(node.cause) 

481 elif hasattr(node, "type") and node.type is not None: 

482 self.visit(node.type) 

483 if node.inst is not None: 

484 self.write(", ") 

485 self.visit(node.inst) 

486 if node.tback is not None: 

487 self.write(", ") 

488 self.visit(node.tback) 

489 

490 # Expressions 

491 

492 def visit_Attribute(self, node): 

493 self.visit(node.value) 

494 self.write("." + node.attr) 

495 

496 def visit_Call(self, node): 

497 want_comma = [] 

498 

499 def write_comma(): 

500 if want_comma: 

501 self.write(", ") 

502 else: 

503 want_comma.append(True) 

504 

505 self.visit(node.func) 

506 self.write("(") 

507 for arg in node.args: 

508 write_comma() 

509 self.visit(arg) 

510 for keyword in node.keywords: 

511 write_comma() 

512 self.write(keyword.arg + "=") 

513 self.visit(keyword.value) 

514 if getattr(node, "starargs", None): 

515 write_comma() 

516 self.write("*") 

517 self.visit(node.starargs) 

518 if getattr(node, "kwargs", None): 

519 write_comma() 

520 self.write("**") 

521 self.visit(node.kwargs) 

522 self.write(")") 

523 

524 def visit_Name(self, node): 

525 self.write(node.id) 

526 

527 def visit_NameConstant(self, node): 

528 self.write(str(node.value)) 

529 

530 def visit_arg(self, node): 

531 self.write(node.arg) 

532 

533 def visit_Str(self, node): 

534 self.write(repr(node.s)) 

535 

536 def visit_Bytes(self, node): 

537 self.write(repr(node.s)) 

538 

539 def visit_Num(self, node): 

540 self.write(repr(node.n)) 

541 

542 # newly needed in Python 3.8 

543 def visit_Constant(self, node): 

544 self.write(repr(node.value)) 

545 

546 def visit_Tuple(self, node): 

547 self.write("(") 

548 idx = -1 

549 for idx, item in enumerate(node.elts): 

550 if idx: 

551 self.write(", ") 

552 self.visit(item) 

553 self.write(idx and ")" or ",)") 

554 

555 def sequence_visit(left, right): 

556 def visit(self, node): 

557 self.write(left) 

558 for idx, item in enumerate(node.elts): 

559 if idx: 

560 self.write(", ") 

561 self.visit(item) 

562 self.write(right) 

563 

564 return visit 

565 

566 visit_List = sequence_visit("[", "]") 

567 visit_Set = sequence_visit("{", "}") 

568 del sequence_visit 

569 

570 def visit_Dict(self, node): 

571 self.write("{") 

572 for idx, (key, value) in enumerate(zip(node.keys, node.values)): 

573 if idx: 

574 self.write(", ") 

575 self.visit(key) 

576 self.write(": ") 

577 self.visit(value) 

578 self.write("}") 

579 

580 def visit_BinOp(self, node): 

581 self.write("(") 

582 self.visit(node.left) 

583 self.write(" %s " % BINOP_SYMBOLS[type(node.op)]) 

584 self.visit(node.right) 

585 self.write(")") 

586 

587 def visit_BoolOp(self, node): 

588 self.write("(") 

589 for idx, value in enumerate(node.values): 

590 if idx: 

591 self.write(" %s " % BOOLOP_SYMBOLS[type(node.op)]) 

592 self.visit(value) 

593 self.write(")") 

594 

595 def visit_Compare(self, node): 

596 self.write("(") 

597 self.visit(node.left) 

598 for op, right in zip(node.ops, node.comparators): 

599 self.write(" %s " % CMPOP_SYMBOLS[type(op)]) 

600 self.visit(right) 

601 self.write(")") 

602 

603 def visit_UnaryOp(self, node): 

604 self.write("(") 

605 op = UNARYOP_SYMBOLS[type(node.op)] 

606 self.write(op) 

607 if op == "not": 

608 self.write(" ") 

609 self.visit(node.operand) 

610 self.write(")") 

611 

612 def visit_Subscript(self, node): 

613 self.visit(node.value) 

614 self.write("[") 

615 self.visit(node.slice) 

616 self.write("]") 

617 

618 def visit_Slice(self, node): 

619 if node.lower is not None: 

620 self.visit(node.lower) 

621 self.write(":") 

622 if node.upper is not None: 

623 self.visit(node.upper) 

624 if node.step is not None: 

625 self.write(":") 

626 if not (isinstance(node.step, Name) and node.step.id == "None"): 

627 self.visit(node.step) 

628 

629 def visit_ExtSlice(self, node): 

630 for idx, item in node.dims: 

631 if idx: 

632 self.write(", ") 

633 self.visit(item) 

634 

635 def visit_Yield(self, node): 

636 self.write("yield ") 

637 self.visit(node.value) 

638 

639 def visit_Lambda(self, node): 

640 self.write("lambda ") 

641 self.signature(node.args) 

642 self.write(": ") 

643 self.visit(node.body) 

644 

645 def visit_Ellipsis(self, node): 

646 self.write("Ellipsis") 

647 

648 def generator_visit(left, right): 

649 def visit(self, node): 

650 self.write(left) 

651 self.visit(node.elt) 

652 for comprehension in node.generators: 

653 self.visit(comprehension) 

654 self.write(right) 

655 

656 return visit 

657 

658 visit_ListComp = generator_visit("[", "]") 

659 visit_GeneratorExp = generator_visit("(", ")") 

660 visit_SetComp = generator_visit("{", "}") 

661 del generator_visit 

662 

663 def visit_DictComp(self, node): 

664 self.write("{") 

665 self.visit(node.key) 

666 self.write(": ") 

667 self.visit(node.value) 

668 for comprehension in node.generators: 

669 self.visit(comprehension) 

670 self.write("}") 

671 

672 def visit_IfExp(self, node): 

673 self.visit(node.body) 

674 self.write(" if ") 

675 self.visit(node.test) 

676 self.write(" else ") 

677 self.visit(node.orelse) 

678 

679 def visit_Starred(self, node): 

680 self.write("*") 

681 self.visit(node.value) 

682 

683 def visit_Repr(self, node): 

684 # XXX: python 2.6 only 

685 self.write("`") 

686 self.visit(node.value) 

687 self.write("`") 

688 

689 # Helper Nodes 

690 

691 def visit_alias(self, node): 

692 self.write(node.name) 

693 if node.asname is not None: 

694 self.write(" as " + node.asname) 

695 

696 def visit_comprehension(self, node): 

697 self.write(" for ") 

698 self.visit(node.target) 

699 self.write(" in ") 

700 self.visit(node.iter) 

701 if node.ifs: 

702 for if_ in node.ifs: 

703 self.write(" if ") 

704 self.visit(if_) 

705 

706 def visit_excepthandler(self, node): 

707 self.newline() 

708 self.write("except") 

709 if node.type is not None: 

710 self.write(" ") 

711 self.visit(node.type) 

712 if node.name is not None: 

713 self.write(" as ") 

714 self.visit(node.name) 

715 self.write(":") 

716 self.body(node.body)