/Lib/lib2to3/fixes/fix_set_literal.py

http://unladen-swallow.googlecode.com/ · Python · 52 lines · 32 code · 11 blank · 9 comment · 4 complexity · 9529989dd16e172f7e7938d49e162a9a MD5 · raw file

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