Merge pull request #6 from scnerd/cleanup_tests

Cleanup tests

tests/test_collapse_literals.py

@@ -30,43 +30,49 @@ def f(y):
         self.assertEqual(f(-1), deco_f(-1))
 
     def test_basic(self):
-        @pragma.collapse_literals(return_source=True)
+        @pragma.collapse_literals
         def f():
             return 1 + 1
 
-        result = dedent('''
+        result = '''
         def f():
             return 2
-        ''')
-        self.assertEqual(f.strip(), result.strip())
+        '''
+
+        self.assertSourceEqual(f, result)
+        self.assertEqual(f(), 2)
 
     def test_vars(self):
-        @pragma.collapse_literals(return_source=True)
+        @pragma.collapse_literals
         def f():
             x = 3
             y = 2
             return x + y
 
-        result = dedent('''
+        result = '''
         def f():
             x = 3
             y = 2
             return 5
-        ''')
-        self.assertEqual(f.strip(), result.strip())
+        '''
+
+        self.assertSourceEqual(f, result)
+        self.assertEqual(f(), 5)
 
     def test_partial(self):
-        @pragma.collapse_literals(return_source=True)
+        @pragma.collapse_literals
         def f(y):
             x = 3
             return x + 2 + y
 
-        result = dedent('''
+        result = '''
         def f(y):
             x = 3
             return 5 + y
-        ''')
-        self.assertEqual(f.strip(), result.strip())
+        '''
+
+        self.assertSourceEqual(f, result)
+        self.assertEqual(f(5), 10)
 
     def test_constant_index(self):
         @pragma.collapse_literals
@@ -79,23 +85,25 @@ def f():
             x = [1, 2, 3]
             return 1
         '''
+
         self.assertSourceEqual(f, result)
         self.assertEqual(f(), 1)
 
     def test_with_unroll(self):
-        @pragma.collapse_literals(return_source=True)
+        @pragma.collapse_literals
         @pragma.unroll
         def f():
             for i in range(3):
                 print(i + 2)
 
-        result = dedent('''
+        result = '''
         def f():
             print(2)
             print(3)
             print(4)
-        ''')
-        self.assertEqual(f.strip(), result.strip())
+        '''
+
+        self.assertSourceEqual(f, result)
 
     # # TODO: Figure out variable levels of specificity...
     # def test_with_objects(self):
@@ -170,20 +178,22 @@ def f(y):
     #     self.assertEqual(f.strip(), result.strip())
 
     def test_constant_conditional_erasure(self):
-        @pragma.collapse_literals(return_source=True)
-        def f(y):
+        @pragma.collapse_literals
+        def f():
             x = 0
             if x <= 0:
                 x = 1
             return x
 
-        result = dedent('''
-        def f(y):
+        result = '''
+        def f():
             x = 0
             x = 1
             return 1
-        ''')
-        self.assertEqual(f.strip(), result.strip())
+        '''
+
+        self.assertSourceEqual(f, result)
+        self.assertEqual(f(), 1)
 
         def fn():
             if x == 0:
@@ -209,31 +219,36 @@ def fn():
             x = 'c'
             return 'c'
         ''')
-        self.assertEqual(pragma.collapse_literals(return_source=True, x=0)(fn).strip(), result0.strip())
-        self.assertEqual(pragma.collapse_literals(return_source=True, x=1)(fn).strip(), result1.strip())
-        self.assertEqual(pragma.collapse_literals(return_source=True, x=2)(fn).strip(), result2.strip())
+
+        self.assertSourceEqual(pragma.collapse_literals(x=0)(fn), result0)
+        self.assertSourceEqual(pragma.collapse_literals(x=1)(fn), result1)
+        self.assertSourceEqual(pragma.collapse_literals(x=2)(fn), result2)
 
     def test_unary(self):
-        @pragma.collapse_literals(return_source=True)
+        @pragma.collapse_literals
         def f():
             return 1 + -5
 
-        result = dedent('''
+        result = '''
         def f():
             return -4
-        ''')
-        self.assertEqual(f.strip(), result.strip())
+        '''
+
+        self.assertSourceEqual(f, result)
+        self.assertEqual(f(), -4)
 
     def test_funcs(self):
-        @pragma.collapse_literals(return_source=True)
+        @pragma.collapse_literals
         def f():
             return sum(range(5))
 
-        result = dedent('''
+        result = '''
         def f():
             return 10
-        ''')
-        self.assertEqual(f.strip(), result.strip())
+        '''
+
+        self.assertSourceEqual(f, result)
+        self.assertEqual(f(), 10)
 
     def test_funcs2(self):
         my_list = [1, 2, 3]
@@ -272,12 +287,13 @@ def f():
             ((y, x), z) = ((1, 2), 3)
             return x
 
-        result = dedent('''
+        result = '''
         def f():
             x = 3
             (y, x), z = (1, 2), 3
             return 2
-        ''')
+        '''
+
         self.assertSourceEqual(f, result)
         self.assertEqual(f(), 2)
 

tests/test_deindex.py

@@ -8,6 +8,7 @@
 class TestDeindex(PragmaTest):
     def test_with_literals(self):
         v = [1, 2, 3]
+
         @pragma.collapse_literals
         @pragma.deindex(v, 'v')
         def f():
@@ -23,6 +24,7 @@ def f():
 
     def test_with_objects(self):
         v = [object(), object(), object()]
+
         @pragma.deindex(v, 'v')
         def f():
             return v[0] + v[1] + v[2]
@@ -34,13 +36,28 @@ def f():
 
         self.assertSourceEqual(f, result)
 
+    def test_with_objects_same_instance(self):
+        v = [object(), object(), object()]
+
+        @pragma.deindex(v, 'v')
+        def f():
+            return v[0]
+
+        result = '''
+        def f():
+            return v_0
+        '''
+
+        self.assertSourceEqual(f, result)
+        self.assertIs(f(), v[0])
+
     def test_with_unroll(self):
         v = [None, None, None]
 
         @pragma.deindex(v, 'v', return_source=True)
-        @pragma.unroll(lv=len(v))
+        @pragma.unroll
         def f():
-            for i in range(lv):
+            for i in range(len(v)):
                 yield v[i]
 
         result = dedent('''
@@ -49,29 +66,24 @@ def f():
             yield v_1
             yield v_2
         ''')
-        self.assertEqual(f.strip(), result.strip())
 
-    def test_with_objects_run(self):
-        v = [object(), object(), object()]
-        @pragma.deindex(v, 'v')
-        def f():
-            return v[0]
-
-        self.assertEqual(f(), v[0])
+        self.assertSourceEqual(f, result)
 
     def test_with_variable_indices(self):
         v = [object(), object(), object()]
-        @pragma.deindex(v, 'v', return_source=True)
+
+        @pragma.deindex(v, 'v')
         def f(x):
             yield v[0]
             yield v[x]
 
-        result = dedent('''
+        result = '''
         def f(x):
             yield v_0
             yield v[x]
-        ''')
-        self.assertEqual(f.strip(), result.strip())
+        '''
+
+        self.assertSourceEqual(f, result)
 
     def test_dict(self):
         d = {'a': 1, 'b': 2}
@@ -108,16 +120,17 @@ def run_func(i, x):
         self.assertEqual(run_func(1, 5), 25)
         self.assertEqual(run_func(2, 5), 125)
 
-        result = dedent('''
+        result = '''
         def run_func(i, x):
             if i == 0:
                 return funcs_0(x)
             if i == 1:
                 return funcs_1(x)
             if i == 2:
                 return funcs_2(x)
-        ''')
-        self.assertEqual(inspect.getsource(run_func).strip(), result.strip())
+        '''
+
+        self.assertSourceEqual(run_func, result)
 
     def test_len(self):
         a = ['a', 'b', 'c']