ast/codegen.py

18 months ago

author
Georg Brandl <georg@python.org>
date
Mon Aug 23 17:07:30 2010 +0000
changeset 354
98ce1ce17c7c
parent 305
d36dea149ae7
permissions
-rw-r--r--

Remove unnecessary step.

     1 # -*- coding: utf-8 -*-
     2 """
     3     codegen
     4     ~~~~~~~
     6     Extension to ast that allow ast -> python code generation.
     8     :copyright: Copyright 2008 by Armin Ronacher.
     9     :license: BSD.
    10 """
    11 from ast import *
    14 def to_source(node, indent_with=' ' * 4, add_line_information=False):
    15     """This function can convert a node tree back into python sourcecode.
    16     This is useful for debugging purposes, especially if you're dealing with
    17     custom asts not generated by python itself.
    19     It could be that the sourcecode is evaluable when the AST itself is not
    20     compilable / evaluable.  The reason for this is that the AST contains some
    21     more data than regular sourcecode does, which is dropped during
    22     conversion.
    24     Each level of indentation is replaced with `indent_with`.  Per default this
    25     parameter is equal to four spaces as suggested by PEP 8, but it might be
    26     adjusted to match the application's styleguide.
    28     If `add_line_information` is set to `True` comments for the line numbers
    29     of the nodes are added to the output.  This can be used to spot wrong line
    30     number information of statement nodes.
    31     """
    32     generator = SourceGenerator(indent_with, add_line_information)
    33     generator.visit(node)
    34     return ''.join(generator.result)
    37 class SourceGenerator(NodeVisitor):
    38     """This visitor is able to transform a well formed syntax tree into python
    39     sourcecode.  For more details have a look at the docstring of the
    40     `node_to_source` function.
    41     """
    43     def __init__(self, indent_with, add_line_information=False):
    44         self.result = []
    45         self.indent_with = indent_with
    46         self.add_line_information = add_line_information
    47         self.indentation = 0
    48         self.new_lines = 0
    50     def write(self, x):
    51         if self.new_lines:
    52             if self.result:
    53                 self.result.append('\n' * self.new_lines)
    54             self.result.append(self.indent_with * self.indentation)
    55             self.new_lines = 0
    56         self.result.append(x)
    58     def newline(self, node=None, extra=0):
    59         self.new_lines = max(self.new_lines, 1 + extra)
    60         if node is not None and self.add_line_information:
    61             self.write('# line: %s' % node.lineno)
    62             self.new_lines = 1
    64     def body(self, statements):
    65         self.new_line = True
    66         self.indentation += 1
    67         for stmt in statements:
    68             self.visit(stmt)
    69         self.indentation -= 1
    71     def body_or_else(self, node):
    72         self.body(node.body)
    73         if node.orelse:
    74             self.newline()
    75             self.write('else:')
    76             self.body(node.orelse)
    78     def signature(self, node):
    79         want_comma = []
    80         def write_comma():
    81             if want_comma:
    82                 self.write(', ')
    83             else:
    84                 want_comma.append(True)
    86         padding = [None] * (len(node.args) - len(node.defaults))
    87         for arg, default in zip(node.args, padding + node.defaults):
    88             write_comma()
    89             self.visit(arg)
    90             if default is not None:
    91                 self.write('=')
    92                 self.visit(default)
    93         if node.vararg is not None:
    94             write_comma()
    95             self.write('*' + node.vararg)
    96         if node.kwarg is not None:
    97             write_comma()
    98             self.write('**' + node.kwarg)
   100     def decorators(self, node):
   101         for decorator in node.decorator_list:
   102             self.newline(decorator)
   103             self.write('@')
   104             self.visit(decorator)
   106     # Statements
   108     def visit_Assign(self, node):
   109         self.newline(node)
   110         for idx, target in enumerate(node.targets):
   111             if idx:
   112                 self.write(', ')
   113             self.visit(target)
   114         self.write(' = ')
   115         self.visit(node.value)
   117     def visit_AugAssign(self, node):
   118         self.newline(node)
   119         self.visit(node.target)
   120         self.write(BINOP_SYMBOLS[type(node.op)] + '=')
   121         self.visit(node.value)
   123     def visit_ImportFrom(self, node):
   124         self.newline(node)
   125         self.write('from %s%s import ' % ('.' * node.level, node.module))
   126         for idx, item in enumerate(node.names):
   127             if idx:
   128                 self.write(', ')
   129             self.write(item)
   131     def visit_Import(self, node):
   132         self.newline(node)
   133         for item in node.names:
   134             self.write('import ')
   135             self.visit(item)
   137     def visit_Expr(self, node):
   138         self.newline(node)
   139         self.generic_visit(node)
   141     def visit_FunctionDef(self, node):
   142         self.newline(extra=1)
   143         self.decorators(node)
   144         self.newline(node)
   145         self.write('def %s(' % node.name)
   146         self.signature(node.args)
   147         self.write('):')
   148         self.body(node.body)
   150     def visit_ClassDef(self, node):
   151         have_args = []
   152         def paren_or_comma():
   153             if have_args:
   154                 self.write(', ')
   155             else:
   156                 have_args.append(True)
   157                 self.write('(')
   159         self.newline(extra=2)
   160         self.decorators(node)
   161         self.newline(node)
   162         self.write('class %s' % node.name)
   163         for base in node.bases:
   164             paren_or_comma()
   165             self.visit(base)
   166         # XXX: the if here is used to keep this module compatible
   167         #      with python 2.6.
   168         if hasattr(node, 'keywords'):
   169             for keyword in node.keywords:
   170                 paren_or_comma()
   171                 self.write(keyword.arg + '=')
   172                 self.visit(keyword.value)
   173             if node.starargs is not None:
   174                 paren_or_comma()
   175                 self.write('*')
   176                 self.visit(node.starargs)
   177             if node.kwargs is not None:
   178                 paren_or_comma()
   179                 self.write('**')
   180                 self.visit(node.kwargs)
   181         self.write(have_args and '):' or ':')
   182         self.body(node.body)
   184     def visit_If(self, node):
   185         self.newline(node)
   186         self.write('if ')
   187         self.visit(node.test)
   188         self.write(':')
   189         self.body(node.body)
   190         while True:
   191             else_ = node.orelse
   192             if len(else_) == 1 and isinstance(else_[0], If):
   193                 node = else_[0]
   194                 self.newline()
   195                 self.write('elif ')
   196                 self.visit(node.test)
   197                 self.write(':')
   198                 self.body(node.body)
   199             else:
   200                 self.newline()
   201                 self.write('else:')
   202                 self.body(else_)
   203                 break
   205     def visit_For(self, node):
   206         self.newline(node)
   207         self.write('for ')
   208         self.visit(node.target)
   209         self.write(' in ')
   210         self.visit(node.iter)
   211         self.write(':')
   212         self.body_or_else(node)
   214     def visit_While(self, node):
   215         self.newline(node)
   216         self.write('while ')
   217         self.visit(node.test)
   218         self.write(':')
   219         self.body_or_else(node)
   221     def visit_With(self, node):
   222         self.newline(node)
   223         self.write('with ')
   224         self.visit(node.context_expr)
   225         if node.optional_vars is not None:
   226             self.write(' as ')
   227             self.visit(node.optional_vars)
   228         self.write(':')
   229         self.body(node.body)
   231     def visit_Pass(self, node):
   232         self.newline(node)
   233         self.write('pass')
   235     def visit_Print(self, node):
   236         # XXX: python 2.6 only
   237         self.newline(node)
   238         self.write('print ')
   239         want_comma = False
   240         if node.dest is not None:
   241             self.write(' >> ')
   242             self.visit(node.dest)
   243             want_comma = True
   244         for value in node.values:
   245             if want_comma:
   246                 self.write(', ')
   247             self.visit(value)
   248             want_comma = True
   249         if not node.nl:
   250             self.write(',')
   252     def visit_Delete(self, node):
   253         self.newline(node)
   254         self.write('del ')
   255         for idx, target in enumerate(node):
   256             if idx:
   257                 self.write(', ')
   258             self.visit(target)
   260     def visit_TryExcept(self, node):
   261         self.newline(node)
   262         self.write('try:')
   263         self.body(node.body)
   264         for handler in node.handlers:
   265             self.visit(handler)
   267     def visit_TryFinally(self, node):
   268         self.newline(node)
   269         self.write('try:')
   270         self.body(node.body)
   271         self.newline(node)
   272         self.write('finally:')
   273         self.body(node.finalbody)
   275     def visit_Global(self, node):
   276         self.newline(node)
   277         self.write('global ' + ', '.join(node.names))
   279     def visit_Nonlocal(self, node):
   280         self.newline(node)
   281         self.write('nonlocal ' + ', '.join(node.names))
   283     def visit_Return(self, node):
   284         self.newline(node)
   285         self.write('return ')
   286         self.visit(node.value)
   288     def visit_Break(self, node):
   289         self.newline(node)
   290         self.write('break')
   292     def visit_Continue(self, node):
   293         self.newline(node)
   294         self.write('continue')
   296     def visit_Raise(self, node):
   297         # XXX: Python 2.6 / 3.0 compatibility
   298         self.newline(node)
   299         self.write('raise')
   300         if hasattr(node, 'exc') and node.exc is not None:
   301             self.write(' ')
   302             self.visit(node.exc)
   303             if node.cause is not None:
   304                 self.write(' from ')
   305                 self.visit(node.cause)
   306         elif hasattr(node, 'type') and node.type is not None:
   307             self.visit(node.type)
   308             if node.inst is not None:
   309                 self.write(', ')
   310                 self.visit(node.inst)
   311             if node.tback is not None:
   312                 self.write(', ')
   313                 self.visit(node.tback)
   315     # Expressions
   317     def visit_Attribute(self, node):
   318         self.visit(node.value)
   319         self.write('.' + node.attr)
   321     def visit_Call(self, node):
   322         want_comma = []
   323         def write_comma():
   324             if want_comma:
   325                 self.write(', ')
   326             else:
   327                 want_comma.append(True)
   329         self.visit(node.func)
   330         self.write('(')
   331         for arg in node.args:
   332             write_comma()
   333             self.visit(arg)
   334         for keyword in node.keywords:
   335             write_comma()
   336             self.write(keyword.arg + '=')
   337             self.visit(keyword.value)
   338         if node.starargs is not None:
   339             write_comma()
   340             self.write('*')
   341             self.visit(node.starargs)
   342         if node.kwargs is not None:
   343             write_comma()
   344             self.write('**')
   345             self.visit(node.kwargs)
   346         self.write(')')
   348     def visit_Name(self, node):
   349         self.write(node.id)
   351     def visit_Str(self, node):
   352         self.write(repr(node.s))
   354     def visit_Bytes(self, node):
   355         self.write(repr(node.s))
   357     def visit_Num(self, node):
   358         self.write(repr(node.n))
   360     def visit_Tuple(self, node):
   361         self.write('(')
   362         idx = -1
   363         for idx, item in enumerate(node.elts):
   364             if idx:
   365                 self.write(', ')
   366             self.visit(item)
   367         self.write(idx and ')' or ',)')
   369     def sequence_visit(left, right):
   370         def visit(self, node):
   371             self.write(left)
   372             for idx, item in enumerate(node.elts):
   373                 if idx:
   374                     self.write(', ')
   375                 self.visit(item)
   376             self.write(right)
   377         return visit
   379     visit_List = sequence_visit('[', ']')
   380     visit_Set = sequence_visit('{', '}')
   381     del sequence_visit
   383     def visit_Dict(self, node):
   384         self.write('{')
   385         for idx, (key, value) in enumerate(zip(node.keys, node.values)):
   386             if idx:
   387                 self.write(', ')
   388             self.visit(key)
   389             self.write(': ')
   390             self.visit(value)
   391         self.write('}')
   393     def visit_BinOp(self, node):
   394         self.visit(node.left)
   395         self.write(' %s ' % BINOP_SYMBOLS[type(node.op)])
   396         self.visit(node.right)
   398     def visit_BoolOp(self, node):
   399         self.write('(')
   400         for idx, value in enumerate(node.values):
   401             if idx:
   402                 self.write(' %s ' % BOOLOP_SYMBOLS[type(node.op)])
   403             self.visit(value)
   404         self.write(')')
   406     def visit_Compare(self, node):
   407         self.write('(')
   408         self.write(node.left)
   409         for op, right in zip(node.ops, node.comparators):
   410             self.write(' %s %%' % CMPOP_SYMBOLS[type(op)])
   411             self.visit(right)
   412         self.write(')')
   414     def visit_UnaryOp(self, node):
   415         self.write('(')
   416         op = UNARYOP_SYMBOLS[type(node.op)]
   417         self.write(op)
   418         if op == 'not':
   419             self.write(' ')
   420         self.visit(node.operand)
   421         self.write(')')
   423     def visit_Subscript(self, node):
   424         self.visit(node.value)
   425         self.write('[')
   426         self.visit(node.slice)
   427         self.write(']')
   429     def visit_Slice(self, node):
   430         if node.lower is not None:
   431             self.visit(node.lower)
   432         self.write(':')
   433         if node.upper is not None:
   434             self.visit(node.upper)
   435         if node.step is not None:
   436             self.write(':')
   437             if not (isinstance(node.step, Name) and node.step.id == 'None'):
   438                 self.visit(node.step)
   440     def visit_ExtSlice(self, node):
   441         for idx, item in node.dims:
   442             if idx:
   443                 self.write(', ')
   444             self.visit(item)
   446     def visit_Yield(self, node):
   447         self.write('yield ')
   448         self.visit(node.value)
   450     def visit_Lambda(self, node):
   451         self.write('lambda ')
   452         self.signature(node.args)
   453         self.write(': ')
   454         self.visit(node.body)
   456     def visit_Ellipsis(self, node):
   457         self.write('Ellipsis')
   459     def generator_visit(left, right):
   460         def visit(self, node):
   461             self.write(left)
   462             self.visit(node.elt)
   463             for comprehension in node.generators:
   464                 self.visit(comprehension)
   465             self.write(right)
   466         return visit
   468     visit_ListComp = generator_visit('[', ']')
   469     visit_GeneratorExp = generator_visit('(', ')')
   470     visit_SetComp = generator_visit('{', '}')
   471     del generator_visit
   473     def visit_DictComp(self, node):
   474         self.write('{')
   475         self.visit(node.key)
   476         self.write(': ')
   477         self.visit(node.value)
   478         for comprehension in node.generators:
   479             self.visit(comprehension)
   480         self.write('}')
   482     def visit_IfExp(self, node):
   483         self.visit(node.body)
   484         self.write(' if ')
   485         self.visit(node.test)
   486         self.write(' else ')
   487         self.visit(node.orelse)
   489     def visit_Starred(self, node):
   490         self.write('*')
   491         self.visit(node.value)
   493     def visit_Repr(self, node):
   494         # XXX: python 2.6 only
   495         self.write('`')
   496         self.visit(node.value)
   497         self.write('`')
   499     # Helper Nodes
   501     def visit_alias(self, node):
   502         self.write(node.name)
   503         if node.asname is not None:
   504             self.write(' as ' + node.asname)
   506     def visit_comprehension(self, node):
   507         self.write(' for ')
   508         self.visit(node.target)
   509         self.write(' in ')
   510         self.visit(node.iter)
   511         if node.ifs:
   512             for if_ in node.ifs:
   513                 self.write(' if ')
   514                 self.visit(if_)
   516     def visit_excepthandler(self, node):
   517         self.newline(node)
   518         self.write('except')
   519         if node.type is not None:
   520             self.write(' ')
   521             self.visit(node.type)
   522             if node.name is not None:
   523                 self.write(' as ')
   524                 self.visit(node.name)
   525         self.write(':')
   526         self.body(node.body)

mercurial