fix_print.py 2.8 KB

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