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    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, "{")]
39        literal.extend(n.clone() for n in items.children)
40        literal.append(pytree.Leaf(token.RBRACE, "}"))
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