fix_set_literal.py 1.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. """
  2. Optional fixer to transform set() calls to set literals.
  3. """
  4. # Author: Benjamin Peterson
  5. from lib2to3 import fixer_base, pytree
  6. from lib2to3.fixer_util import token, syms
  7. class FixSetLiteral(fixer_base.BaseFix):
  8. BM_compatible = True
  9. explicit = True
  10. PATTERN = """power< 'set' trailer< '('
  11. (atom=atom< '[' (items=listmaker< any ((',' any)* [',']) >
  12. |
  13. single=any) ']' >
  14. |
  15. atom< '(' items=testlist_gexp< any ((',' any)* [',']) > ')' >
  16. )
  17. ')' > >
  18. """
  19. def transform(self, node, results):
  20. single = results.get("single")
  21. if single:
  22. # Make a fake listmaker
  23. fake = pytree.Node(syms.listmaker, [single.clone()])
  24. single.replace(fake)
  25. items = fake
  26. else:
  27. items = results["items"]
  28. # Build the contents of the literal
  29. literal = [pytree.Leaf(token.LBRACE, "{")]
  30. literal.extend(n.clone() for n in items.children)
  31. literal.append(pytree.Leaf(token.RBRACE, "}"))
  32. # Set the prefix of the right brace to that of the ')' or ']'
  33. literal[-1].prefix = items.next_sibling.prefix
  34. maker = pytree.Node(syms.dictsetmaker, literal)
  35. maker.prefix = node.prefix
  36. # If the original was a one tuple, we need to remove the extra comma.
  37. if len(maker.children) == 4:
  38. n = maker.children[2]
  39. n.remove()
  40. maker.children[-1].prefix = n.prefix
  41. # Finally, replace the set call with our shiny new literal.
  42. return maker