| """ | |
| Optional fixer to transform set() calls to set literals. | |
| """ | |
| # Author: Benjamin Peterson | |
| from lib2to3 import fixer_base, pytree | |
| from lib2to3.fixer_util import token, syms | |
| class FixSetLiteral(fixer_base.BaseFix): | |
| BM_compatible = True | |
| explicit = True | |
| PATTERN = """power< 'set' trailer< '(' | |
| (atom=atom< '[' (items=listmaker< any ((',' any)* [',']) > | |
| | | |
| single=any) ']' > | |
| | | |
| atom< '(' items=testlist_gexp< any ((',' any)* [',']) > ')' > | |
| ) | |
| ')' > > | |
| """ | |
| def transform(self, node, results): | |
| single = results.get("single") | |
| if single: | |
| # Make a fake listmaker | |
| fake = pytree.Node(syms.listmaker, [single.clone()]) | |
| single.replace(fake) | |
| items = fake | |
| else: | |
| items = results["items"] | |
| # Build the contents of the literal | |
| literal = [pytree.Leaf(token.LBRACE, u"{")] | |
| literal.extend(n.clone() for n in items.children) | |
| literal.append(pytree.Leaf(token.RBRACE, u"}")) | |
| # Set the prefix of the right brace to that of the ')' or ']' | |
| literal[-1].prefix = items.next_sibling.prefix | |
| maker = pytree.Node(syms.dictsetmaker, literal) | |
| maker.prefix = node.prefix | |
| # If the original was a one tuple, we need to remove the extra comma. | |
| if len(maker.children) == 4: | |
| n = maker.children[2] | |
| n.remove() | |
| maker.children[-1].prefix = n.prefix | |
| # Finally, replace the set call with our shiny new literal. | |
| return maker |