10c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi""" 20c5958b1636c47ed7c284f859c8e805fd06a0e6Bill YiOptional fixer to transform set() calls to set literals. 30c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi""" 40c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 50c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi# Author: Benjamin Peterson 60c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 70c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yifrom lib2to3 import fixer_base, pytree 80c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yifrom lib2to3.fixer_util import token, syms 90c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass FixSetLiteral(fixer_base.BaseFix): 130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi BM_compatible = True 150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi explicit = True 160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi PATTERN = """power< 'set' trailer< '(' 180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi (atom=atom< '[' (items=listmaker< any ((',' any)* [',']) > 190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi | 200c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi single=any) ']' > 210c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi | 220c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi atom< '(' items=testlist_gexp< any ((',' any)* [',']) > ')' > 230c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi ) 240c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi ')' > > 250c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi """ 260c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 270c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def transform(self, node, results): 280c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi single = results.get("single") 290c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if single: 300c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Make a fake listmaker 310c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi fake = pytree.Node(syms.listmaker, [single.clone()]) 320c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi single.replace(fake) 330c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi items = fake 340c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi else: 350c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi items = results["items"] 360c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 370c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Build the contents of the literal 380c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi literal = [pytree.Leaf(token.LBRACE, u"{")] 390c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi literal.extend(n.clone() for n in items.children) 400c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi literal.append(pytree.Leaf(token.RBRACE, u"}")) 410c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Set the prefix of the right brace to that of the ')' or ']' 420c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi literal[-1].prefix = items.next_sibling.prefix 430c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi maker = pytree.Node(syms.dictsetmaker, literal) 440c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi maker.prefix = node.prefix 450c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 460c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # If the original was a one tuple, we need to remove the extra comma. 470c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if len(maker.children) == 4: 480c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi n = maker.children[2] 490c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi n.remove() 500c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi maker.children[-1].prefix = n.prefix 510c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 520c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Finally, replace the set call with our shiny new literal. 530c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return maker 54