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