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# -*- coding: utf-8 -*- 

2# 

3# Copyright (C) 2008-2009 Edgewall Software 

4# All rights reserved. 

5# 

6# This software is licensed as described in the file COPYING, which 

7# you should have received as part of this distribution. The terms 

8# are also available at http://genshi.edgewall.org/wiki/License. 

9# 

10# This software consists of voluntary contributions made by many 

11# individuals. For the exact contribution history, see the revision 

12# history and logs, available at http://genshi.edgewall.org/log/. 

13 

14"""Support classes for generating code from abstract syntax trees.""" 

15 

16try: 

17 import ast 

18except ImportError: 

19 from chameleon import ast25 as ast 

20 

21import sys 

22import logging 

23import weakref 

24import collections 

25 

26 

27AST_NONE = ast.Name(id='None', ctx=ast.Load()) 

28 

29node_annotations = weakref.WeakKeyDictionary() 

30 

31try: 

32 node_annotations[ast.Name()] = None 

33except TypeError: 

34 logging.debug( 

35 "Unable to create weak references to AST nodes. " \ 

36 "A lock will be used around compilation loop." 

37 ) 

38 

39 node_annotations = {} 

40 

41__docformat__ = 'restructuredtext en' 

42 

43 

44def annotated(value): 

45 node = load("annotation") 

46 node_annotations[node] = value 

47 return node 

48 

49 

50def parse(source, mode='eval'): 

51 return compile(source, '', mode, ast.PyCF_ONLY_AST) 

52 

53 

54def load(name): 

55 return ast.Name(id=name, ctx=ast.Load()) 

56 

57 

58def store(name): 

59 return ast.Name(id=name, ctx=ast.Store()) 

60 

61 

62def param(name): 

63 return ast.Name(id=name, ctx=ast.Param()) 

64 

65 

66def delete(name): 

67 return ast.Name(id=name, ctx=ast.Del()) 

68 

69 

70def subscript(name, value, ctx): 

71 return ast.Subscript( 

72 value=value, 

73 slice=ast.Index(value=ast.Str(s=name)), 

74 ctx=ctx, 

75 ) 

76 

77 

78def walk_names(target, mode): 

79 for node in ast.walk(target): 

80 if isinstance(node, ast.Name) and \ 

81 isinstance(node.ctx, mode): 

82 yield node.id 

83 

84 

85def iter_fields(node): 

86 """ 

87 Yield a tuple of ``(fieldname, value)`` for each field in ``node._fields`` 

88 that is present on *node*. 

89 """ 

90 for field in node._fields: 

91 try: 

92 yield field, getattr(node, field) 

93 except AttributeError: 

94 pass 

95 

96 

97def iter_child_nodes(node): 

98 """ 

99 Yield all direct child nodes of *node*, that is, all fields that are nodes 

100 and all items of fields that are lists of nodes. 

101 """ 

102 for name, field in iter_fields(node): 

103 if isinstance(field, Node): 

104 yield field 

105 elif isinstance(field, list): 

106 for item in field: 

107 if isinstance(item, Node): 

108 yield item 

109 

110 

111def walk(node): 

112 """ 

113 Recursively yield all descendant nodes in the tree starting at *node* 

114 (including *node* itself), in no specified order. This is useful if you 

115 only want to modify nodes in place and don't care about the context. 

116 """ 

117 todo = collections.deque([node]) 

118 while todo: 

119 node = todo.popleft() 

120 todo.extend(iter_child_nodes(node)) 

121 yield node 

122 

123 

124def copy(source, target): 

125 target.__class__ = source.__class__ 

126 target.__dict__ = source.__dict__ 

127 

128 

129def swap(body, replacement, name): 

130 root = ast.Expression(body=body) 

131 for node in ast.walk(root): 

132 if (isinstance(node, ast.Name) and 

133 isinstance(node.ctx, ast.Load) and 

134 node.id == name): 

135 assert hasattr(replacement, '_fields') 

136 node_annotations.setdefault(node, replacement) 

137 

138 

139def marker(name): 

140 return ast.Str(s="__%s" % name) 

141 

142 

143class Node(object): 

144 """AST baseclass that gives us a convenient initialization 

145 method. We explicitly declare and use the ``_fields`` attribute.""" 

146 

147 _fields = () 

148 

149 def __init__(self, *args, **kwargs): 

150 assert isinstance(self._fields, tuple) 

151 self.__dict__.update(kwargs) 

152 for name, value in zip(self._fields, args): 

153 setattr(self, name, value) 

154 

155 def __repr__(self): 

156 """Poor man's single-line pretty printer.""" 

157 

158 name = type(self).__name__ 

159 return '<%s%s at %x>' % ( 

160 name, 

161 "".join(" %s=%r" % (name, getattr(self, name, "\"?\"")) 

162 for name in self._fields), 

163 id(self) 

164 ) 

165 

166 def extract(self, condition): 

167 result = [] 

168 for node in walk(self): 

169 if condition(node): 

170 result.append(node) 

171 

172 return result 

173 

174 

175class Builtin(Node): 

176 """Represents a Python builtin. 

177 

178 Used when a builtin is used internally by the compiler, to avoid 

179 clashing with a user assignment (e.g. ``help`` is a builtin, but 

180 also commonly assigned in templates). 

181 """ 

182 

183 _fields = "id", "ctx" 

184 

185 ctx = ast.Load() 

186 

187 

188class Symbol(Node): 

189 """Represents an importable symbol.""" 

190 

191 _fields = "value", 

192 

193 

194class Static(Node): 

195 """Represents a static value.""" 

196 

197 _fields = "value", "name" 

198 

199 name = None 

200 

201 

202class Comment(Node): 

203 _fields = "text", "space", "stmt" 

204 

205 stmt = None 

206 space = "" 

207 

208 

209class TokenRef(Node): 

210 """Represents a source-code token reference.""" 

211 

212 _fields = "pos", "length" 

213 

214 

215class ASTCodeGenerator(object): 

216 """General purpose base class for AST transformations. 

217 

218 Every visitor method can be overridden to return an AST node that has been 

219 altered or replaced in some way. 

220 """ 

221 

222 def __init__(self, tree): 

223 self.lines_info = [] 

224 self.line_info = [] 

225 self.lines = [] 

226 self.line = "" 

227 self.last = None 

228 self.indent = 0 

229 self.blame_stack = [] 

230 self.visit(tree) 

231 

232 if self.line.strip(): 

233 self._new_line() 

234 

235 self.line = None 

236 self.line_info = None 

237 

238 # strip trivial lines 

239 self.code = "\n".join( 

240 line.strip() and line or "" 

241 for line in self.lines 

242 ) 

243 

244 def _change_indent(self, delta): 

245 self.indent += delta 

246 

247 def _new_line(self): 

248 if self.line is not None: 

249 self.lines.append(self.line) 

250 self.lines_info.append(self.line_info) 

251 self.line = ' ' * 4 * self.indent 

252 if len(self.blame_stack) == 0: 

253 self.line_info = [] 

254 self.last = None 

255 else: 

256 self.line_info = [(0, self.blame_stack[-1],)] 

257 self.last = self.blame_stack[-1] 

258 

259 def _write(self, s): 

260 if len(s) == 0: 

261 return 

262 if len(self.blame_stack) == 0: 

263 if self.last is not None: 

264 self.last = None 

265 self.line_info.append((len(self.line), self.last)) 

266 else: 

267 if self.last != self.blame_stack[-1]: 

268 self.last = self.blame_stack[-1] 

269 self.line_info.append((len(self.line), self.last)) 

270 self.line += s 

271 

272 def flush(self): 

273 if self.line: 

274 self._new_line() 

275 

276 def visit(self, node): 

277 if node is None: 

278 return None 

279 if type(node) is tuple: 

280 return tuple([self.visit(n) for n in node]) 

281 try: 

282 self.blame_stack.append((node.lineno, node.col_offset,)) 

283 info = True 

284 except AttributeError: 

285 info = False 

286 visitor = getattr(self, 'visit_%s' % node.__class__.__name__, None) 

287 if visitor is None: 

288 raise Exception('No handler for ``%s`` (%s).' % ( 

289 node.__class__.__name__, repr(node))) 

290 ret = visitor(node) 

291 if info: 

292 self.blame_stack.pop() 

293 return ret 

294 

295 def visit_Module(self, node): 

296 for n in node.body: 

297 self.visit(n) 

298 visit_Interactive = visit_Module 

299 visit_Suite = visit_Module 

300 

301 def visit_Expression(self, node): 

302 return self.visit(node.body) 

303 

304 # arguments = (expr* args, identifier? vararg, 

305 # identifier? kwarg, expr* defaults) 

306 def visit_arguments(self, node): 

307 first = True 

308 no_default_count = len(node.args) - len(node.defaults) 

309 for i, arg in enumerate(node.args): 

310 if not first: 

311 self._write(', ') 

312 else: 

313 first = False 

314 self.visit(arg) 

315 if i >= no_default_count: 

316 self._write('=') 

317 self.visit(node.defaults[i - no_default_count]) 

318 if getattr(node, 'vararg', None): 

319 if not first: 

320 self._write(', ') 

321 else: 

322 first = False 

323 self._write('*' + node.vararg) 

324 if getattr(node, 'kwarg', None): 

325 if not first: 

326 self._write(', ') 

327 else: 

328 first = False 

329 self._write('**' + node.kwarg) 

330 

331 def visit_arg(self, node): 

332 self._write(node.arg) 

333 

334 # FunctionDef(identifier name, arguments args, 

335 # stmt* body, expr* decorators) 

336 def visit_FunctionDef(self, node): 

337 self._new_line() 

338 for decorator in getattr(node, 'decorator_list', ()): 

339 self._new_line() 

340 self._write('@') 

341 self.visit(decorator) 

342 self._new_line() 

343 self._write('def ' + node.name + '(') 

344 self.visit(node.args) 

345 self._write('):') 

346 self._change_indent(1) 

347 for statement in node.body: 

348 self.visit(statement) 

349 self._change_indent(-1) 

350 

351 # ClassDef(identifier name, expr* bases, stmt* body) 

352 def visit_ClassDef(self, node): 

353 self._new_line() 

354 self._write('class ' + node.name) 

355 if node.bases: 

356 self._write('(') 

357 self.visit(node.bases[0]) 

358 for base in node.bases[1:]: 

359 self._write(', ') 

360 self.visit(base) 

361 self._write(')') 

362 self._write(':') 

363 self._change_indent(1) 

364 for statement in node.body: 

365 self.visit(statement) 

366 self._change_indent(-1) 

367 

368 # Return(expr? value) 

369 def visit_Return(self, node): 

370 self._new_line() 

371 self._write('return') 

372 if getattr(node, 'value', None): 

373 self._write(' ') 

374 self.visit(node.value) 

375 

376 # Delete(expr* targets) 

377 def visit_Delete(self, node): 

378 self._new_line() 

379 self._write('del ') 

380 self.visit(node.targets[0]) 

381 for target in node.targets[1:]: 

382 self._write(', ') 

383 self.visit(target) 

384 

385 # Assign(expr* targets, expr value) 

386 def visit_Assign(self, node): 

387 self._new_line() 

388 for target in node.targets: 

389 self.visit(target) 

390 self._write(' = ') 

391 self.visit(node.value) 

392 

393 # AugAssign(expr target, operator op, expr value) 

394 def visit_AugAssign(self, node): 

395 self._new_line() 

396 self.visit(node.target) 

397 self._write(' ' + self.binary_operators[node.op.__class__] + '= ') 

398 self.visit(node.value) 

399 

400 # JoinedStr(expr* values) 

401 def visit_JoinedStr(self, node): 

402 if node.values: 

403 self._write('"".join((') 

404 for value in node.values: 

405 self.visit(value) 

406 self._write(',') 

407 self._write('))') 

408 else: 

409 self._write('""') 

410 

411 # FormattedValue(expr value) 

412 def visit_FormattedValue(self, node): 

413 if node.conversion == ord('r'): 

414 self._write('repr') 

415 elif node.conversion == ord('a'): 

416 self._write('ascii') 

417 else: 

418 self._write('str') 

419 self._write('(') 

420 self.visit(node.value) 

421 if node.format_spec is not None: 

422 self._write(').__format__(') 

423 self.visit(node.format_spec) 

424 self._write(')') 

425 

426 # Print(expr? dest, expr* values, bool nl) 

427 def visit_Print(self, node): 

428 self._new_line() 

429 self._write('print') 

430 if getattr(node, 'dest', None): 

431 self._write(' >> ') 

432 self.visit(node.dest) 

433 if getattr(node, 'values', None): 

434 self._write(', ') 

435 else: 

436 self._write(' ') 

437 if getattr(node, 'values', None): 

438 self.visit(node.values[0]) 

439 for value in node.values[1:]: 

440 self._write(', ') 

441 self.visit(value) 

442 if not node.nl: 

443 self._write(',') 

444 

445 # For(expr target, expr iter, stmt* body, stmt* orelse) 

446 def visit_For(self, node): 

447 self._new_line() 

448 self._write('for ') 

449 self.visit(node.target) 

450 self._write(' in ') 

451 self.visit(node.iter) 

452 self._write(':') 

453 self._change_indent(1) 

454 for statement in node.body: 

455 self.visit(statement) 

456 self._change_indent(-1) 

457 if getattr(node, 'orelse', None): 

458 self._new_line() 

459 self._write('else:') 

460 self._change_indent(1) 

461 for statement in node.orelse: 

462 self.visit(statement) 

463 self._change_indent(-1) 

464 

465 # While(expr test, stmt* body, stmt* orelse) 

466 def visit_While(self, node): 

467 self._new_line() 

468 self._write('while ') 

469 self.visit(node.test) 

470 self._write(':') 

471 self._change_indent(1) 

472 for statement in node.body: 

473 self.visit(statement) 

474 self._change_indent(-1) 

475 if getattr(node, 'orelse', None): 

476 self._new_line() 

477 self._write('else:') 

478 self._change_indent(1) 

479 for statement in node.orelse: 

480 self.visit(statement) 

481 self._change_indent(-1) 

482 

483 # If(expr test, stmt* body, stmt* orelse) 

484 def visit_If(self, node): 

485 self._new_line() 

486 self._write('if ') 

487 self.visit(node.test) 

488 self._write(':') 

489 self._change_indent(1) 

490 for statement in node.body: 

491 self.visit(statement) 

492 self._change_indent(-1) 

493 if getattr(node, 'orelse', None): 

494 self._new_line() 

495 self._write('else:') 

496 self._change_indent(1) 

497 for statement in node.orelse: 

498 self.visit(statement) 

499 self._change_indent(-1) 

500 

501 # With(expr context_expr, expr? optional_vars, stmt* body) 

502 def visit_With(self, node): 

503 self._new_line() 

504 self._write('with ') 

505 self.visit(node.context_expr) 

506 if getattr(node, 'optional_vars', None): 

507 self._write(' as ') 

508 self.visit(node.optional_vars) 

509 self._write(':') 

510 self._change_indent(1) 

511 for statement in node.body: 

512 self.visit(statement) 

513 self._change_indent(-1) 

514 

515 # Raise(expr? type, expr? inst, expr? tback) 

516 def visit_Raise(self, node): 

517 self._new_line() 

518 self._write('raise') 

519 if not getattr(node, "type", None): 

520 exc = getattr(node, "exc", None) 

521 if exc is None: 

522 return 

523 self._write(' ') 

524 return self.visit(exc) 

525 self._write(' ') 

526 self.visit(node.type) 

527 if not node.inst: 

528 return 

529 self._write(', ') 

530 self.visit(node.inst) 

531 if not node.tback: 

532 return 

533 self._write(', ') 

534 self.visit(node.tback) 

535 

536 # Try(stmt* body, excepthandler* handlers, stmt* orelse, stmt* finalbody) 

537 def visit_Try(self, node): 

538 self._new_line() 

539 self._write('try:') 

540 self._change_indent(1) 

541 for statement in node.body: 

542 self.visit(statement) 

543 self._change_indent(-1) 

544 if getattr(node, 'handlers', None): 

545 for handler in node.handlers: 

546 self.visit(handler) 

547 self._new_line() 

548 

549 if getattr(node, 'orelse', None): 

550 self._write('else:') 

551 self._change_indent(1) 

552 for statement in node.orelse: 

553 self.visit(statement) 

554 self._change_indent(-1) 

555 

556 if getattr(node, 'finalbody', None): 

557 self._new_line() 

558 self._write('finally:') 

559 self._change_indent(1) 

560 for statement in node.finalbody: 

561 self.visit(statement) 

562 self._change_indent(-1) 

563 

564 # TryExcept(stmt* body, excepthandler* handlers, stmt* orelse) 

565 def visit_TryExcept(self, node): 

566 self._new_line() 

567 self._write('try:') 

568 self._change_indent(1) 

569 for statement in node.body: 

570 self.visit(statement) 

571 self._change_indent(-1) 

572 if getattr(node, 'handlers', None): 

573 for handler in node.handlers: 

574 self.visit(handler) 

575 self._new_line() 

576 if getattr(node, 'orelse', None): 

577 self._write('else:') 

578 self._change_indent(1) 

579 for statement in node.orelse: 

580 self.visit(statement) 

581 self._change_indent(-1) 

582 

583 # excepthandler = (expr? type, expr? name, stmt* body) 

584 def visit_ExceptHandler(self, node): 

585 self._new_line() 

586 self._write('except') 

587 if getattr(node, 'type', None): 

588 self._write(' ') 

589 self.visit(node.type) 

590 if getattr(node, 'name', None): 

591 if sys.version_info[0] == 2: 

592 assert getattr(node, 'type', None) 

593 self._write(', ') 

594 else: 

595 self._write(' as ') 

596 self.visit(node.name) 

597 self._write(':') 

598 self._change_indent(1) 

599 for statement in node.body: 

600 self.visit(statement) 

601 self._change_indent(-1) 

602 visit_excepthandler = visit_ExceptHandler 

603 

604 # TryFinally(stmt* body, stmt* finalbody) 

605 def visit_TryFinally(self, node): 

606 self._new_line() 

607 self._write('try:') 

608 self._change_indent(1) 

609 for statement in node.body: 

610 self.visit(statement) 

611 self._change_indent(-1) 

612 

613 if getattr(node, 'finalbody', None): 

614 self._new_line() 

615 self._write('finally:') 

616 self._change_indent(1) 

617 for statement in node.finalbody: 

618 self.visit(statement) 

619 self._change_indent(-1) 

620 

621 # Assert(expr test, expr? msg) 

622 def visit_Assert(self, node): 

623 self._new_line() 

624 self._write('assert ') 

625 self.visit(node.test) 

626 if getattr(node, 'msg', None): 

627 self._write(', ') 

628 self.visit(node.msg) 

629 

630 def visit_alias(self, node): 

631 self._write(node.name) 

632 if getattr(node, 'asname', None): 

633 self._write(' as ') 

634 self._write(node.asname) 

635 

636 # Import(alias* names) 

637 def visit_Import(self, node): 

638 self._new_line() 

639 self._write('import ') 

640 self.visit(node.names[0]) 

641 for name in node.names[1:]: 

642 self._write(', ') 

643 self.visit(name) 

644 

645 # ImportFrom(identifier module, alias* names, int? level) 

646 def visit_ImportFrom(self, node): 

647 self._new_line() 

648 self._write('from ') 

649 if node.level: 

650 self._write('.' * node.level) 

651 self._write(node.module) 

652 self._write(' import ') 

653 self.visit(node.names[0]) 

654 for name in node.names[1:]: 

655 self._write(', ') 

656 self.visit(name) 

657 

658 # Exec(expr body, expr? globals, expr? locals) 

659 def visit_Exec(self, node): 

660 self._new_line() 

661 self._write('exec ') 

662 self.visit(node.body) 

663 if not node.globals: 

664 return 

665 self._write(', ') 

666 self.visit(node.globals) 

667 if not node.locals: 

668 return 

669 self._write(', ') 

670 self.visit(node.locals) 

671 

672 # Global(identifier* names) 

673 def visit_Global(self, node): 

674 self._new_line() 

675 self._write('global ') 

676 self.visit(node.names[0]) 

677 for name in node.names[1:]: 

678 self._write(', ') 

679 self.visit(name) 

680 

681 # Expr(expr value) 

682 def visit_Expr(self, node): 

683 self._new_line() 

684 self.visit(node.value) 

685 

686 # Pass 

687 def visit_Pass(self, node): 

688 self._new_line() 

689 self._write('pass') 

690 

691 # Break 

692 def visit_Break(self, node): 

693 self._new_line() 

694 self._write('break') 

695 

696 # Continue 

697 def visit_Continue(self, node): 

698 self._new_line() 

699 self._write('continue') 

700 

701 ### EXPRESSIONS 

702 def with_parens(f): 

703 def _f(self, node): 

704 self._write('(') 

705 f(self, node) 

706 self._write(')') 

707 return _f 

708 

709 bool_operators = {ast.And: 'and', ast.Or: 'or'} 

710 

711 # BoolOp(boolop op, expr* values) 

712 @with_parens 

713 def visit_BoolOp(self, node): 

714 joiner = ' ' + self.bool_operators[node.op.__class__] + ' ' 

715 self.visit(node.values[0]) 

716 for value in node.values[1:]: 

717 self._write(joiner) 

718 self.visit(value) 

719 

720 binary_operators = { 

721 ast.Add: '+', 

722 ast.Sub: '-', 

723 ast.Mult: '*', 

724 ast.Div: '/', 

725 ast.Mod: '%', 

726 ast.Pow: '**', 

727 ast.LShift: '<<', 

728 ast.RShift: '>>', 

729 ast.BitOr: '|', 

730 ast.BitXor: '^', 

731 ast.BitAnd: '&', 

732 ast.FloorDiv: '//' 

733 } 

734 

735 # BinOp(expr left, operator op, expr right) 

736 @with_parens 

737 def visit_BinOp(self, node): 

738 self.visit(node.left) 

739 self._write(' ' + self.binary_operators[node.op.__class__] + ' ') 

740 self.visit(node.right) 

741 

742 unary_operators = { 

743 ast.Invert: '~', 

744 ast.Not: 'not', 

745 ast.UAdd: '+', 

746 ast.USub: '-', 

747 } 

748 

749 # UnaryOp(unaryop op, expr operand) 

750 def visit_UnaryOp(self, node): 

751 self._write(self.unary_operators[node.op.__class__] + ' ') 

752 self.visit(node.operand) 

753 

754 # Lambda(arguments args, expr body) 

755 @with_parens 

756 def visit_Lambda(self, node): 

757 self._write('lambda ') 

758 self.visit(node.args) 

759 self._write(': ') 

760 self.visit(node.body) 

761 

762 # IfExp(expr test, expr body, expr orelse) 

763 @with_parens 

764 def visit_IfExp(self, node): 

765 self.visit(node.body) 

766 self._write(' if ') 

767 self.visit(node.test) 

768 self._write(' else ') 

769 self.visit(node.orelse) 

770 

771 # Dict(expr* keys, expr* values) 

772 def visit_Dict(self, node): 

773 self._write('{') 

774 for key, value in zip(node.keys, node.values): 

775 self.visit(key) 

776 self._write(': ') 

777 self.visit(value) 

778 self._write(', ') 

779 self._write('}') 

780 

781 def visit_Set(self, node): 

782 self._write('{') 

783 elts = list(node.elts) 

784 last = elts.pop() 

785 for elt in elts: 

786 self.visit(elt) 

787 self._write(', ') 

788 self.visit(last) 

789 self._write('}') 

790 

791 # ListComp(expr elt, comprehension* generators) 

792 def visit_ListComp(self, node): 

793 self._write('[') 

794 self.visit(node.elt) 

795 for generator in node.generators: 

796 # comprehension = (expr target, expr iter, expr* ifs) 

797 self._write(' for ') 

798 self.visit(generator.target) 

799 self._write(' in ') 

800 self.visit(generator.iter) 

801 for ifexpr in generator.ifs: 

802 self._write(' if ') 

803 self.visit(ifexpr) 

804 self._write(']') 

805 

806 # GeneratorExp(expr elt, comprehension* generators) 

807 def visit_GeneratorExp(self, node): 

808 self._write('(') 

809 self.visit(node.elt) 

810 for generator in node.generators: 

811 # comprehension = (expr target, expr iter, expr* ifs) 

812 self._write(' for ') 

813 self.visit(generator.target) 

814 self._write(' in ') 

815 self.visit(generator.iter) 

816 for ifexpr in generator.ifs: 

817 self._write(' if ') 

818 self.visit(ifexpr) 

819 self._write(')') 

820 

821 # Yield(expr? value) 

822 def visit_Yield(self, node): 

823 self._write('yield') 

824 if getattr(node, 'value', None): 

825 self._write(' ') 

826 self.visit(node.value) 

827 

828 comparison_operators = { 

829 ast.Eq: '==', 

830 ast.NotEq: '!=', 

831 ast.Lt: '<', 

832 ast.LtE: '<=', 

833 ast.Gt: '>', 

834 ast.GtE: '>=', 

835 ast.Is: 'is', 

836 ast.IsNot: 'is not', 

837 ast.In: 'in', 

838 ast.NotIn: 'not in', 

839 } 

840 

841 # Compare(expr left, cmpop* ops, expr* comparators) 

842 @with_parens 

843 def visit_Compare(self, node): 

844 self.visit(node.left) 

845 for op, comparator in zip(node.ops, node.comparators): 

846 self._write(' ' + self.comparison_operators[op.__class__] + ' ') 

847 self.visit(comparator) 

848 

849 # Call(expr func, expr* args, keyword* keywords, 

850 # expr? starargs, expr? kwargs) 

851 def visit_Call(self, node): 

852 self.visit(node.func) 

853 self._write('(') 

854 first = True 

855 for arg in node.args: 

856 if not first: 

857 self._write(', ') 

858 first = False 

859 self.visit(arg) 

860 

861 for keyword in node.keywords: 

862 if not first: 

863 self._write(', ') 

864 first = False 

865 # keyword = (identifier arg, expr value) 

866 if keyword.arg is not None: 

867 self._write(keyword.arg) 

868 self._write('=') 

869 else: 

870 self._write('**') 

871 self.visit(keyword.value) 

872 # Attribute removed in Python 3.5 

873 if getattr(node, 'starargs', None): 

874 if not first: 

875 self._write(', ') 

876 first = False 

877 self._write('*') 

878 self.visit(node.starargs) 

879 

880 # Attribute removed in Python 3.5 

881 if getattr(node, 'kwargs', None): 

882 if not first: 

883 self._write(', ') 

884 first = False 

885 self._write('**') 

886 self.visit(node.kwargs) 

887 self._write(')') 

888 

889 # Repr(expr value) 

890 def visit_Repr(self, node): 

891 self._write('`') 

892 self.visit(node.value) 

893 self._write('`') 

894 

895 # Constant(object value) 

896 def visit_Constant(self, node): 

897 if node.value is Ellipsis: 

898 self._write('...') 

899 else: 

900 self._write(repr(node.value)) 

901 

902 # Num(object n) 

903 def visit_Num(self, node): 

904 self._write(repr(node.n)) 

905 

906 # Str(string s) 

907 def visit_Str(self, node): 

908 self._write(repr(node.s)) 

909 

910 def visit_Ellipsis(self, node): 

911 self._write('...') 

912 

913 # Attribute(expr value, identifier attr, expr_context ctx) 

914 def visit_Attribute(self, node): 

915 self.visit(node.value) 

916 self._write('.') 

917 self._write(node.attr) 

918 

919 # Subscript(expr value, slice slice, expr_context ctx) 

920 def visit_Subscript(self, node): 

921 self.visit(node.value) 

922 self._write('[') 

923 if isinstance(node.slice, ast.Tuple) and node.slice.elts: 

924 self.visit(node.slice.elts[0]) 

925 if len(node.slice.elts) == 1: 

926 self._write(', ') 

927 else: 

928 for dim in node.slice.elts[1:]: 

929 self._write(', ') 

930 self.visit(dim) 

931 elif isinstance(node.slice, ast.Slice): 

932 self.visit_Slice(node.slice, True) 

933 else: 

934 self.visit(node.slice) 

935 self._write(']') 

936 

937 # Slice(expr? lower, expr? upper, expr? step) 

938 def visit_Slice(self, node, subscription=False): 

939 if subscription: 

940 if getattr(node, 'lower', None) is not None: 

941 self.visit(node.lower) 

942 self._write(':') 

943 if getattr(node, 'upper', None) is not None: 

944 self.visit(node.upper) 

945 if getattr(node, 'step', None) is not None: 

946 self._write(':') 

947 self.visit(node.step) 

948 else: 

949 self._write('slice(') 

950 self.visit(getattr(node, "lower", None) or AST_NONE) 

951 self._write(', ') 

952 self.visit(getattr(node, "upper", None) or AST_NONE) 

953 self._write(', ') 

954 self.visit(getattr(node, "step", None) or AST_NONE) 

955 self._write(')') 

956 

957 # Index(expr value) 

958 def visit_Index(self, node): 

959 self.visit(node.value) 

960 

961 # ExtSlice(slice* dims) 

962 def visit_ExtSlice(self, node): 

963 self.visit(node.dims[0]) 

964 if len(node.dims) == 1: 

965 self._write(', ') 

966 else: 

967 for dim in node.dims[1:]: 

968 self._write(', ') 

969 self.visit(dim) 

970 

971 # Starred(expr value, expr_context ctx) 

972 def visit_Starred(self, node): 

973 self._write('*') 

974 self.visit(node.value) 

975 

976 # Name(identifier id, expr_context ctx) 

977 def visit_Name(self, node): 

978 self._write(node.id) 

979 

980 # List(expr* elts, expr_context ctx) 

981 def visit_List(self, node): 

982 self._write('[') 

983 for elt in node.elts: 

984 self.visit(elt) 

985 self._write(', ') 

986 self._write(']') 

987 

988 # Tuple(expr *elts, expr_context ctx) 

989 def visit_Tuple(self, node): 

990 self._write('(') 

991 for elt in node.elts: 

992 self.visit(elt) 

993 self._write(', ') 

994 self._write(')') 

995 

996 # NameConstant(singleton value) 

997 def visit_NameConstant(self, node): 

998 self._write(str(node.value)) 

999 

1000class AnnotationAwareVisitor(ast.NodeVisitor): 

1001 def visit(self, node): 

1002 annotation = node_annotations.get(node) 

1003 if annotation is not None: 

1004 assert hasattr(annotation, '_fields') 

1005 node = annotation 

1006 

1007 super(AnnotationAwareVisitor, self).visit(node) 

1008 

1009 def apply_transform(self, node): 

1010 if node not in node_annotations: 

1011 result = self.transform(node) 

1012 if result is not None and result is not node: 

1013 node_annotations[node] = result 

1014 

1015 

1016class NameLookupRewriteVisitor(AnnotationAwareVisitor): 

1017 def __init__(self, transform): 

1018 self.transform = transform 

1019 self.transformed = set() 

1020 self.scopes = [set()] 

1021 

1022 def __call__(self, node): 

1023 self.visit(node) 

1024 return self.transformed 

1025 

1026 def visit_arg(self, node): 

1027 scope = self.scopes[-1] 

1028 scope.add(node.arg) 

1029 

1030 def visit_Name(self, node): 

1031 scope = self.scopes[-1] 

1032 if isinstance(node.ctx, ast.Param): 

1033 scope.add(node.id) 

1034 elif node.id not in scope: 

1035 self.transformed.add(node.id) 

1036 self.apply_transform(node) 

1037 

1038 def visit_FunctionDef(self, node): 

1039 self.scopes[-1].add(node.name) 

1040 

1041 def visit_alias(self, node): 

1042 name = node.asname if node.asname is not None else node.name 

1043 self.scopes[-1].add(name) 

1044 

1045 def visit_Lambda(self, node): 

1046 self.scopes.append(set()) 

1047 try: 

1048 self.visit(node.args) 

1049 self.visit(node.body) 

1050 finally: 

1051 self.scopes.pop() 

1052 

1053 

1054class ItemLookupOnAttributeErrorVisitor(AnnotationAwareVisitor): 

1055 def __init__(self, transform): 

1056 self.transform = transform 

1057 

1058 def visit_Attribute(self, node): 

1059 self.generic_visit(node) 

1060 self.apply_transform(node)