/Lib/lib2to3/fixes/fix_set_literal.py

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