fix_tuple_params.py 5.4 KB

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