[Python-checkins] r66893 - in sandbox/trunk/2to3/lib2to3: fixes/fix_set_literal.py tests/test_fixers.py

[benjamin.peterson]

Author: benjamin.peterson
Date: Wed Oct 15 00:16:54 2008
New Revision: 66893

Log:
add an optional set literal fixer

Added:
   sandbox/trunk/2to3/lib2to3/fixes/fix_set_literal.py   (contents, props changed)
Modified:
   sandbox/trunk/2to3/lib2to3/tests/test_fixers.py

Added: sandbox/trunk/2to3/lib2to3/fixes/fix_set_literal.py
==============================================================================
--- (empty file)
+++ sandbox/trunk/2to3/lib2to3/fixes/fix_set_literal.py	Wed Oct 15 00:16:54 2008
@@ -0,0 +1,52 @@
+"""
+Optional fixer to transform set() calls to set literals.
+"""
+
+# Author: Benjamin Peterson
+
+from lib2to3 import fixer_base, pytree
+from lib2to3.fixer_util import token, syms
+
+
+
+class FixSetLiteral(fixer_base.BaseFix):
+
+    explicit = True
+
+    PATTERN = """power< 'set' trailer< '('
+                     (atom=atom< '[' (items=listmaker< any ((',' any)* [',']) >
+                                |
+                                single=any) ']' >
+                     |
+                     atom< '(' items=testlist_gexp< any ((',' any)* [',']) > ')' >
+                     )
+                     ')' > >
+              """
+
+    def transform(self, node, results):
+        single = results.get("single")
+        if single:
+            # Make a fake listmaker
+            fake = pytree.Node(syms.listmaker, [single.clone()])
+            single.replace(fake)
+            items = fake
+        else:
+            items = results["items"]
+
+        # Build the contents of the literal
+        literal = [pytree.Leaf(token.LBRACE, "{")]
+        literal.extend(n.clone() for n in items.children)
+        literal.append(pytree.Leaf(token.RBRACE, "}"))
+        # Set the prefix of the right brace to that of the ')' or ']'
+        literal[-1].set_prefix(items.get_next_sibling().get_prefix())
+        maker = pytree.Node(syms.dictsetmaker, literal)
+        maker.set_prefix(node.get_prefix())
+
+        # If the original was a one tuple, we need to remove the extra comma.
+        if len(maker.children) == 4:
+            n = maker.children[2]
+            n.remove()
+            maker.children[-1].set_prefix(n.get_prefix())
+
+        # Finally, replace the set call with our shiny new literal.
+        return maker

Modified: sandbox/trunk/2to3/lib2to3/tests/test_fixers.py
==============================================================================
--- sandbox/trunk/2to3/lib2to3/tests/test_fixers.py	(original)
+++ sandbox/trunk/2to3/lib2to3/tests/test_fixers.py	Wed Oct 15 00:16:54 2008
@@ -3383,6 +3383,134 @@
         """
         self.check_both(b, a)
 
+
+class Test_set_literal(FixerTestCase):
+
+    fixer = "set_literal"
+
+    def test_basic(self):
+        b = """set([1, 2, 3])"""
+        a = """{1, 2, 3}"""
+        self.check(b, a)
+
+        b = """set((1, 2, 3))"""
+        a = """{1, 2, 3}"""
+        self.check(b, a)
+
+        b = """set((1,))"""
+        a = """{1}"""
+        self.check(b, a)
+
+        b = """set([1])"""
+        self.check(b, a)
+
+        b = """set((a, b))"""
+        a = """{a, b}"""
+        self.check(b, a)
+
+        b = """set([a, b])"""
+        self.check(b, a)
+
+        b = """set((a*234, f(args=23)))"""
+        a = """{a*234, f(args=23)}"""
+        self.check(b, a)
+
+        b = """set([a*23, f(23)])"""
+        a = """{a*23, f(23)}"""
+        self.check(b, a)
+
+        b = """set([a-234**23])"""
+        a = """{a-234**23}"""
+        self.check(b, a)
+
+    def test_listcomps(self):
+        b = """set([x for x in y])"""
+        a = """{x for x in y}"""
+        self.check(b, a)
+
+        b = """set([x for x in y if x == m])"""
+        a = """{x for x in y if x == m}"""
+        self.check(b, a)
+
+        b = """set([x for x in y for a in b])"""
+        a = """{x for x in y for a in b}"""
+        self.check(b, a)
+
+        b = """set([f(x) - 23 for x in y])"""
+        a = """{f(x) - 23 for x in y}"""
+        self.check(b, a)
+
+    def test_whitespace(self):
+        b = """set( [1, 2])"""
+        a = """{1, 2}"""
+        self.check(b, a)
+
+        b = """set([1 ,  2])"""
+        a = """{1 ,  2}"""
+        self.check(b, a)
+
+        b = """set([ 1 ])"""
+        a = """{ 1 }"""
+        self.check(b, a)
+
+        b = """set( [1] )"""
+        a = """{1}"""
+        self.check(b, a)
+
+        b = """set([  1,  2  ])"""
+        a = """{  1,  2  }"""
+        self.check(b, a)
+
+        b = """set([x  for x in y ])"""
+        a = """{x  for x in y }"""
+        self.check(b, a)
+
+        b = """set(
+                   [1, 2]
+               )
+            """
+        a = """{1, 2}\n"""
+        self.check(b, a)
+
+    def test_comments(self):
+        b = """set((1, 2)) # Hi"""
+        a = """{1, 2} # Hi"""
+        self.check(b, a)
+
+        # This isn't optimal behavior, but the fixer is optional.
+        b = """
+            # Foo
+            set( # Bar
+               (1, 2)
+            )
+            """
+        a = """
+            # Foo
+            {1, 2}
+            """
+        self.check(b, a)
+
+    def test_unchanged(self):
+        s = """set()"""
+        self.unchanged(s)
+
+        s = """set(a)"""
+        self.unchanged(s)
+
+        s = """set(a, b, c)"""
+        self.unchanged(s)
+
+        # Don't transform generators because they might have to be lazy.
+        s = """set(x for x in y)"""
+        self.unchanged(s)
+
+        s = """set(x for x in y if z)"""
+        self.unchanged(s)
+
+        s = """set(a*823-23**2 + f(23))"""
+        self.unchanged(s)
+
+
 class Test_sys_exc(FixerTestCase):
     fixer = "sys_exc"