fix_unpacking.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124
  1. u"""
  2. Fixer for:
  3. (a,)* *b (,c)* [,] = s
  4. for (a,)* *b (,c)* [,] in d: ...
  5. """
  6. from lib2to3 import fixer_base
  7. from itertools import count
  8. from lib2to3.fixer_util import (Assign, Comma, Call, Newline, Name,
  9. Number, token, syms, Node, Leaf)
  10. from libfuturize.fixer_util import indentation, suitify, commatize
  11. # from libfuturize.fixer_util import Assign, Comma, Call, Newline, Name, Number, indentation, suitify, commatize, token, syms, Node, Leaf
  12. def assignment_source(num_pre, num_post, LISTNAME, ITERNAME):
  13. u"""
  14. Accepts num_pre and num_post, which are counts of values
  15. before and after the starg (not including the starg)
  16. Returns a source fit for Assign() from fixer_util
  17. """
  18. children = []
  19. try:
  20. pre = unicode(num_pre)
  21. post = unicode(num_post)
  22. except NameError:
  23. pre = str(num_pre)
  24. post = str(num_post)
  25. # This code builds the assignment source from lib2to3 tree primitives.
  26. # It's not very readable, but it seems like the most correct way to do it.
  27. if num_pre > 0:
  28. pre_part = Node(syms.power, [Name(LISTNAME), Node(syms.trailer, [Leaf(token.LSQB, u"["), Node(syms.subscript, [Leaf(token.COLON, u":"), Number(pre)]), Leaf(token.RSQB, u"]")])])
  29. children.append(pre_part)
  30. children.append(Leaf(token.PLUS, u"+", prefix=u" "))
  31. main_part = Node(syms.power, [Leaf(token.LSQB, u"[", prefix=u" "), Name(LISTNAME), Node(syms.trailer, [Leaf(token.LSQB, u"["), Node(syms.subscript, [Number(pre) if num_pre > 0 else Leaf(1, u""), Leaf(token.COLON, u":"), Node(syms.factor, [Leaf(token.MINUS, u"-"), Number(post)]) if num_post > 0 else Leaf(1, u"")]), Leaf(token.RSQB, u"]"), Leaf(token.RSQB, u"]")])])
  32. children.append(main_part)
  33. if num_post > 0:
  34. children.append(Leaf(token.PLUS, u"+", prefix=u" "))
  35. post_part = Node(syms.power, [Name(LISTNAME, prefix=u" "), Node(syms.trailer, [Leaf(token.LSQB, u"["), Node(syms.subscript, [Node(syms.factor, [Leaf(token.MINUS, u"-"), Number(post)]), Leaf(token.COLON, u":")]), Leaf(token.RSQB, u"]")])])
  36. children.append(post_part)
  37. source = Node(syms.arith_expr, children)
  38. return source
  39. class FixUnpacking(fixer_base.BaseFix):
  40. PATTERN = u"""
  41. expl=expr_stmt< testlist_star_expr<
  42. pre=(any ',')*
  43. star_expr< '*' name=NAME >
  44. post=(',' any)* [','] > '=' source=any > |
  45. impl=for_stmt< 'for' lst=exprlist<
  46. pre=(any ',')*
  47. star_expr< '*' name=NAME >
  48. post=(',' any)* [','] > 'in' it=any ':' suite=any>"""
  49. def fix_explicit_context(self, node, results):
  50. pre, name, post, source = (results.get(n) for n in (u"pre", u"name", u"post", u"source"))
  51. pre = [n.clone() for n in pre if n.type == token.NAME]
  52. name.prefix = u" "
  53. post = [n.clone() for n in post if n.type == token.NAME]
  54. target = [n.clone() for n in commatize(pre + [name.clone()] + post)]
  55. # to make the special-case fix for "*z, = ..." correct with the least
  56. # amount of modification, make the left-side into a guaranteed tuple
  57. target.append(Comma())
  58. source.prefix = u""
  59. setup_line = Assign(Name(self.LISTNAME), Call(Name(u"list"), [source.clone()]))
  60. power_line = Assign(target, assignment_source(len(pre), len(post), self.LISTNAME, self.ITERNAME))
  61. return setup_line, power_line
  62. def fix_implicit_context(self, node, results):
  63. u"""
  64. Only example of the implicit context is
  65. a for loop, so only fix that.
  66. """
  67. pre, name, post, it = (results.get(n) for n in (u"pre", u"name", u"post", u"it"))
  68. pre = [n.clone() for n in pre if n.type == token.NAME]
  69. name.prefix = u" "
  70. post = [n.clone() for n in post if n.type == token.NAME]
  71. target = [n.clone() for n in commatize(pre + [name.clone()] + post)]
  72. # to make the special-case fix for "*z, = ..." correct with the least
  73. # amount of modification, make the left-side into a guaranteed tuple
  74. target.append(Comma())
  75. source = it.clone()
  76. source.prefix = u""
  77. setup_line = Assign(Name(self.LISTNAME), Call(Name(u"list"), [Name(self.ITERNAME)]))
  78. power_line = Assign(target, assignment_source(len(pre), len(post), self.LISTNAME, self.ITERNAME))
  79. return setup_line, power_line
  80. def transform(self, node, results):
  81. u"""
  82. a,b,c,d,e,f,*g,h,i = range(100) changes to
  83. _3to2list = list(range(100))
  84. a,b,c,d,e,f,g,h,i, = _3to2list[:6] + [_3to2list[6:-2]] + _3to2list[-2:]
  85. and
  86. for a,b,*c,d,e in iter_of_iters: do_stuff changes to
  87. for _3to2iter in iter_of_iters:
  88. _3to2list = list(_3to2iter)
  89. a,b,c,d,e, = _3to2list[:2] + [_3to2list[2:-2]] + _3to2list[-2:]
  90. do_stuff
  91. """
  92. self.LISTNAME = self.new_name(u"_3to2list")
  93. self.ITERNAME = self.new_name(u"_3to2iter")
  94. expl, impl = results.get(u"expl"), results.get(u"impl")
  95. if expl is not None:
  96. setup_line, power_line = self.fix_explicit_context(node, results)
  97. setup_line.prefix = expl.prefix
  98. power_line.prefix = indentation(expl.parent)
  99. setup_line.append_child(Newline())
  100. parent = node.parent
  101. i = node.remove()
  102. parent.insert_child(i, power_line)
  103. parent.insert_child(i, setup_line)
  104. elif impl is not None:
  105. setup_line, power_line = self.fix_implicit_context(node, results)
  106. suitify(node)
  107. suite = [k for k in node.children if k.type == syms.suite][0]
  108. setup_line.prefix = u""
  109. power_line.prefix = suite.children[1].value
  110. suite.children[2].prefix = indentation(suite.children[2])
  111. suite.insert_child(2, Newline())
  112. suite.insert_child(2, power_line)
  113. suite.insert_child(2, Newline())
  114. suite.insert_child(2, setup_line)
  115. results.get(u"lst").replace(Name(self.ITERNAME, prefix=u" "))