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     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