/Lib/lib2to3/fixes/fix_tuple_params.py

http://unladen-swallow.googlecode.com/ · Python · 169 lines · 114 code · 21 blank · 34 comment · 33 complexity · 557690cc5399b0ade14c16089df2effb MD5 · raw file

  1. """Fixer for function definitions with tuple parameters.
  2. def func(((a, b), c), d):
  3. ...
  4. ->
  5. def func(x, d):
  6. ((a, b), c) = x
  7. ...
  8. It will also support lambdas:
  9. lambda (x, y): x + y -> lambda t: t[0] + t[1]
  10. # The parens are a syntax error in Python 3
  11. lambda (x): x + y -> lambda x: x + y
  12. """
  13. # Author: Collin Winter
  14. # Local imports
  15. from .. import pytree
  16. from ..pgen2 import token
  17. from .. import fixer_base
  18. from ..fixer_util import Assign, Name, Newline, Number, Subscript, syms
  19. def is_docstring(stmt):
  20. return isinstance(stmt, pytree.Node) and \
  21. stmt.children[0].type == token.STRING
  22. class FixTupleParams(fixer_base.BaseFix):
  23. PATTERN = """
  24. funcdef< 'def' any parameters< '(' args=any ')' >
  25. ['->' any] ':' suite=any+ >
  26. |
  27. lambda=
  28. lambdef< 'lambda' args=vfpdef< '(' inner=any ')' >
  29. ':' body=any
  30. >
  31. """
  32. def transform(self, node, results):
  33. if "lambda" in results:
  34. return self.transform_lambda(node, results)
  35. new_lines = []
  36. suite = results["suite"]
  37. args = results["args"]
  38. # This crap is so "def foo(...): x = 5; y = 7" is handled correctly.
  39. # TODO(cwinter): suite-cleanup
  40. if suite[0].children[1].type == token.INDENT:
  41. start = 2
  42. indent = suite[0].children[1].value
  43. end = Newline()
  44. else:
  45. start = 0
  46. indent = "; "
  47. end = pytree.Leaf(token.INDENT, "")
  48. # We need access to self for new_name(), and making this a method
  49. # doesn't feel right. Closing over self and new_lines makes the
  50. # code below cleaner.
  51. def handle_tuple(tuple_arg, add_prefix=False):
  52. n = Name(self.new_name())
  53. arg = tuple_arg.clone()
  54. arg.set_prefix("")
  55. stmt = Assign(arg, n.clone())
  56. if add_prefix:
  57. n.set_prefix(" ")
  58. tuple_arg.replace(n)
  59. new_lines.append(pytree.Node(syms.simple_stmt,
  60. [stmt, end.clone()]))
  61. if args.type == syms.tfpdef:
  62. handle_tuple(args)
  63. elif args.type == syms.typedargslist:
  64. for i, arg in enumerate(args.children):
  65. if arg.type == syms.tfpdef:
  66. # Without add_prefix, the emitted code is correct,
  67. # just ugly.
  68. handle_tuple(arg, add_prefix=(i > 0))
  69. if not new_lines:
  70. return node
  71. # This isn't strictly necessary, but it plays nicely with other fixers.
  72. # TODO(cwinter) get rid of this when children becomes a smart list
  73. for line in new_lines:
  74. line.parent = suite[0]
  75. # TODO(cwinter) suite-cleanup
  76. after = start
  77. if start == 0:
  78. new_lines[0].set_prefix(" ")
  79. elif is_docstring(suite[0].children[start]):
  80. new_lines[0].set_prefix(indent)
  81. after = start + 1
  82. suite[0].children[after:after] = new_lines
  83. for i in range(after+1, after+len(new_lines)+1):
  84. suite[0].children[i].set_prefix(indent)
  85. suite[0].changed()
  86. def transform_lambda(self, node, results):
  87. args = results["args"]
  88. body = results["body"]
  89. inner = simplify_args(results["inner"])
  90. # Replace lambda ((((x)))): x with lambda x: x
  91. if inner.type == token.NAME:
  92. inner = inner.clone()
  93. inner.set_prefix(" ")
  94. args.replace(inner)
  95. return
  96. params = find_params(args)
  97. to_index = map_to_index(params)
  98. tup_name = self.new_name(tuple_name(params))
  99. new_param = Name(tup_name, prefix=" ")
  100. args.replace(new_param.clone())
  101. for n in body.post_order():
  102. if n.type == token.NAME and n.value in to_index:
  103. subscripts = [c.clone() for c in to_index[n.value]]
  104. new = pytree.Node(syms.power,
  105. [new_param.clone()] + subscripts)
  106. new.set_prefix(n.get_prefix())
  107. n.replace(new)
  108. ### Helper functions for transform_lambda()
  109. def simplify_args(node):
  110. if node.type in (syms.vfplist, token.NAME):
  111. return node
  112. elif node.type == syms.vfpdef:
  113. # These look like vfpdef< '(' x ')' > where x is NAME
  114. # or another vfpdef instance (leading to recursion).
  115. while node.type == syms.vfpdef:
  116. node = node.children[1]
  117. return node
  118. raise RuntimeError("Received unexpected node %s" % node)
  119. def find_params(node):
  120. if node.type == syms.vfpdef:
  121. return find_params(node.children[1])
  122. elif node.type == token.NAME:
  123. return node.value
  124. return [find_params(c) for c in node.children if c.type != token.COMMA]
  125. def map_to_index(param_list, prefix=[], d=None):
  126. if d is None:
  127. d = {}
  128. for i, obj in enumerate(param_list):
  129. trailer = [Subscript(Number(i))]
  130. if isinstance(obj, list):
  131. map_to_index(obj, trailer, d=d)
  132. else:
  133. d[obj] = prefix + trailer
  134. return d
  135. def tuple_name(param_list):
  136. l = []
  137. for obj in param_list:
  138. if isinstance(obj, list):
  139. l.append(tuple_name(obj))
  140. else:
  141. l.append(obj)
  142. return "_".join(l)