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