Home | History | Annotate | Download | only in fixes 
      1 # Copyright 2006 Google, Inc. All Rights Reserved.
      2 # Licensed to PSF under a Contributor Agreement.
      3 
      4 """Fixer for print.
      5 
      6 Change:
      7     'print'          into 'print()'
      8     'print ...'      into 'print(...)'
      9     'print ... ,'    into 'print(..., end=" ")'
     10     'print >>x, ...' into 'print(..., file=x)'
     11 
     12 No changes are applied if print_function is imported from __future__
     13 
     14 """
     15 
     16 # Local imports
     17 from .. import patcomp
     18 from .. import pytree
     19 from ..pgen2 import token
     20 from .. import fixer_base
     21 from ..fixer_util import Name, Call, Comma, String, is_tuple
     22 
     23 
     24 parend_expr = patcomp.compile_pattern(
     25               """atom< '(' [atom|STRING|NAME] ')' >"""
     26               )
     27 
     28 
     29 class FixPrint(fixer_base.BaseFix):
     30 
     31     BM_compatible = True
     32 
     33     PATTERN = """
     34               simple_stmt< any* bare='print' any* > | print_stmt
     35               """
     36 
     37     def transform(self, node, results):
     38         assert results
     39 
     40         bare_print = results.get("bare")
     41 
     42         if bare_print:
     43             # Special-case print all by itself
     44             bare_print.replace(Call(Name(u"print"), [],
     45                                prefix=bare_print.prefix))
     46             return
     47         assert node.children[0] == Name(u"print")
     48         args = node.children[1:]
     49         if len(args) == 1 and parend_expr.match(args[0]):
     50             # We don't want to keep sticking parens around an
     51             # already-parenthesised expression.
     52             return
     53 
     54         sep = end = file = None
     55         if args and args[-1] == Comma():
     56             args = args[:-1]
     57             end = " "
     58         if args and args[0] == pytree.Leaf(token.RIGHTSHIFT, u">>"):
     59             assert len(args) >= 2
     60             file = args[1].clone()
     61             args = args[3:] # Strip a possible comma after the file expression
     62         # Now synthesize a print(args, sep=..., end=..., file=...) node.
     63         l_args = [arg.clone() for arg in args]
     64         if l_args:
     65             l_args[0].prefix = u""
     66         if sep is not None or end is not None or file is not None:
     67             if sep is not None:
     68                 self.add_kwarg(l_args, u"sep", String(repr(sep)))
     69             if end is not None:
     70                 self.add_kwarg(l_args, u"end", String(repr(end)))
     71             if file is not None:
     72                 self.add_kwarg(l_args, u"file", file)
     73         n_stmt = Call(Name(u"print"), l_args)
     74         n_stmt.prefix = node.prefix
     75         return n_stmt
     76 
     77     def add_kwarg(self, l_nodes, s_kwd, n_expr):
     78         # XXX All this prefix-setting may lose comments (though rarely)
     79         n_expr.prefix = u""
     80         n_argument = pytree.Node(self.syms.argument,
     81                                  (Name(s_kwd),
     82                                   pytree.Leaf(token.EQUAL, u"="),
     83                                   n_expr))
     84         if l_nodes:
     85             l_nodes.append(Comma())
     86             n_argument.prefix = u" "
     87         l_nodes.append(n_argument)
     88