Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/chameleon/astutil.py : 19%

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/.
14"""Support classes for generating code from abstract syntax trees."""
16try:
17 import ast
18except ImportError:
19 from chameleon import ast25 as ast
21import sys
22import logging
23import weakref
24import collections
27AST_NONE = ast.Name(id='None', ctx=ast.Load())
29node_annotations = weakref.WeakKeyDictionary()
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 )
39 node_annotations = {}
41__docformat__ = 'restructuredtext en'
44def annotated(value):
45 node = load("annotation")
46 node_annotations[node] = value
47 return node
50def parse(source, mode='eval'):
51 return compile(source, '', mode, ast.PyCF_ONLY_AST)
54def load(name):
55 return ast.Name(id=name, ctx=ast.Load())
58def store(name):
59 return ast.Name(id=name, ctx=ast.Store())
62def param(name):
63 return ast.Name(id=name, ctx=ast.Param())
66def delete(name):
67 return ast.Name(id=name, ctx=ast.Del())
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 )
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
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
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
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
124def copy(source, target):
125 target.__class__ = source.__class__
126 target.__dict__ = source.__dict__
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)
139def marker(name):
140 return ast.Str(s="__%s" % name)
143class Node(object):
144 """AST baseclass that gives us a convenient initialization
145 method. We explicitly declare and use the ``_fields`` attribute."""
147 _fields = ()
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)
155 def __repr__(self):
156 """Poor man's single-line pretty printer."""
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 )
166 def extract(self, condition):
167 result = []
168 for node in walk(self):
169 if condition(node):
170 result.append(node)
172 return result
175class Builtin(Node):
176 """Represents a Python builtin.
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 """
183 _fields = "id", "ctx"
185 ctx = ast.Load()
188class Symbol(Node):
189 """Represents an importable symbol."""
191 _fields = "value",
194class Static(Node):
195 """Represents a static value."""
197 _fields = "value", "name"
199 name = None
202class Comment(Node):
203 _fields = "text", "space", "stmt"
205 stmt = None
206 space = ""
209class TokenRef(Node):
210 """Represents a source-code token reference."""
212 _fields = "pos", "length"
215class ASTCodeGenerator(object):
216 """General purpose base class for AST transformations.
218 Every visitor method can be overridden to return an AST node that has been
219 altered or replaced in some way.
220 """
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)
232 if self.line.strip():
233 self._new_line()
235 self.line = None
236 self.line_info = None
238 # strip trivial lines
239 self.code = "\n".join(
240 line.strip() and line or ""
241 for line in self.lines
242 )
244 def _change_indent(self, delta):
245 self.indent += delta
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]
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
272 def flush(self):
273 if self.line:
274 self._new_line()
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
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
301 def visit_Expression(self, node):
302 return self.visit(node.body)
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)
331 def visit_arg(self, node):
332 self._write(node.arg)
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)
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)
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)
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)
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)
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)
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('""')
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(')')
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(',')
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)
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)
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)
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)
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)
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()
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)
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)
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)
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
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)
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)
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)
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)
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)
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)
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)
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)
681 # Expr(expr value)
682 def visit_Expr(self, node):
683 self._new_line()
684 self.visit(node.value)
686 # Pass
687 def visit_Pass(self, node):
688 self._new_line()
689 self._write('pass')
691 # Break
692 def visit_Break(self, node):
693 self._new_line()
694 self._write('break')
696 # Continue
697 def visit_Continue(self, node):
698 self._new_line()
699 self._write('continue')
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
709 bool_operators = {ast.And: 'and', ast.Or: 'or'}
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)
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 }
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)
742 unary_operators = {
743 ast.Invert: '~',
744 ast.Not: 'not',
745 ast.UAdd: '+',
746 ast.USub: '-',
747 }
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)
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)
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)
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('}')
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('}')
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(']')
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(')')
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)
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 }
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)
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)
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)
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(')')
889 # Repr(expr value)
890 def visit_Repr(self, node):
891 self._write('`')
892 self.visit(node.value)
893 self._write('`')
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))
902 # Num(object n)
903 def visit_Num(self, node):
904 self._write(repr(node.n))
906 # Str(string s)
907 def visit_Str(self, node):
908 self._write(repr(node.s))
910 def visit_Ellipsis(self, node):
911 self._write('...')
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)
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(']')
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(')')
957 # Index(expr value)
958 def visit_Index(self, node):
959 self.visit(node.value)
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)
971 # Starred(expr value, expr_context ctx)
972 def visit_Starred(self, node):
973 self._write('*')
974 self.visit(node.value)
976 # Name(identifier id, expr_context ctx)
977 def visit_Name(self, node):
978 self._write(node.id)
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(']')
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(')')
996 # NameConstant(singleton value)
997 def visit_NameConstant(self, node):
998 self._write(str(node.value))
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
1007 super(AnnotationAwareVisitor, self).visit(node)
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
1016class NameLookupRewriteVisitor(AnnotationAwareVisitor):
1017 def __init__(self, transform):
1018 self.transform = transform
1019 self.transformed = set()
1020 self.scopes = [set()]
1022 def __call__(self, node):
1023 self.visit(node)
1024 return self.transformed
1026 def visit_arg(self, node):
1027 scope = self.scopes[-1]
1028 scope.add(node.arg)
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)
1038 def visit_FunctionDef(self, node):
1039 self.scopes[-1].add(node.name)
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)
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()
1054class ItemLookupOnAttributeErrorVisitor(AnnotationAwareVisitor):
1055 def __init__(self, transform):
1056 self.transform = transform
1058 def visit_Attribute(self, node):
1059 self.generic_visit(node)
1060 self.apply_transform(node)