Merge pull request #22 from scnerd/21-fix-unroll-issue

21 fix unroll issue

.bumpversion.cfg

@@ -0,0 +1,8 @@
+[bumpversion]
+current_version = 0.2.5
+commit = True
+tag = True
+
+[bumpversion:file:setup.py]
+
+[bumpversion:file:docs/conf.py]

.travis.yml

@@ -4,6 +4,7 @@ python:
  - "3.6"
  - "3.7"
  - "3.8"
+ - "3.9"
 dist: xenial
 
 install:

docs/conf.py

@@ -62,7 +62,7 @@
 # The short X.Y version.
 version = '0.2'
 # The full version, including alpha/beta/rc tags.
-release = '0.2.4'
+release = '0.2.5'
 
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.

docs/todo.rst

@@ -6,5 +6,6 @@ TODO List
 .. todo:: Technically, ``x += y`` doesn't have to be the same thing as ``x = x + y``. Handle it as its own operation of the form ``x += y; return x``
 .. todo:: Support efficiently inlining simple functions, i.e. where there is no return or only one return as the last line of the function, using pure name substitution without loops, try/except, or anything else fancy
 .. todo:: Catch replacement of loop variables that conflict with globals, or throw a more descriptive error when detected. See ``test_iteration_variable``
+.. todo:: Python 3.8/3.9+ support: https://docs.python.org/3/library/ast.html . ``ast.Constant`` is taking over for ``ast.[Num, Str, Bytes, NameConstant, Ellipsis]``. Simple-valued indexes are now values, and extended slices are now tuples: ``ast.[Index, ExtSlice]`` no longer exist.
 
 .. todolist::

pragma/collapse_literals.py

@@ -52,7 +52,8 @@ def resolve_attr_of_slice(attr):
             resolve_attr_of_slice('upper')
             resolve_attr_of_slice('step')
         else:
-            raise TypeError(type(target.slice))
+            target.slice = self.visit(target.slice)  # The index could be anything in 3.8+
+            # raise TypeError(type(target.slice))
 
     def visit_Assign(self, node):
         for it, target in enumerate(node.targets):

pragma/core/transformer.py

@@ -3,7 +3,6 @@
 import inspect
 import logging
 import sys
-import tempfile
 import textwrap
 import warnings
 
@@ -14,6 +13,7 @@
 from .stack import DictStack
 from .resolve import resolve_literal, resolve_iterable, resolve_indexable, resolve_name_or_attribute, \
     make_ast_from_literal
+from ..utils import save_or_return_source
 
 log = logging.getLogger(__name__)
 
@@ -404,34 +404,7 @@ def inner(f):
             trans.unroll_in_tiers = unroll_in_tiers
             f_mod.body[0].decorator_list = []
             f_mod = trans.visit(f_mod)
-            # print(astor.dump_tree(f_mod))
-            if return_source or save_source:
-                try:
-                    source = astor.to_source(f_mod)
-                except Exception as ex:  # pragma: nocover
-                    raise RuntimeError(astor.dump_tree(f_mod)) from ex
-            else:
-                source = None
-
-            if return_source:
-                return source
-            else:
-                f_mod = ast.fix_missing_locations(f_mod)
-                if save_source:
-                    temp = tempfile.NamedTemporaryFile('w', delete=False)
-                    f_file = temp.name
-                exec(compile(f_mod, f_file, 'exec'), glbls)
-                func = glbls[f_mod.body[0].name]
-                if save_source:
-                    func.__tempfile__ = temp
-                    # When there are other decorators, the co_firstlineno of *some* python distributions gets confused
-                    # and thinks they will be there even when they are not written to the file, causing readline overflow
-                    # So we put some empty lines to make them align
-                    temp.write(source)
-                    temp.write('\n' * func.__code__.co_firstlineno)
-                    temp.flush()
-                    temp.close()
-                return func
+            return save_or_return_source(f_file, f_mod, glbls, return_source, save_source)
 
         return inner
 

pragma/lift.py

@@ -9,6 +9,7 @@
 
 from .core.resolve import make_ast_from_literal
 from .core.transformer import function_ast
+from .utils import save_or_return_source
 
 log = logging.getLogger(__name__)
 
@@ -158,30 +159,4 @@ def __call__(self, f):
         )
 
         f_mod.body[0] = new_func_def
-
-        if self.return_source or self.save_source:
-            try:
-                source = astor.to_source(f_mod)
-            except ImportError:  # pragma: nocover
-                raise ImportError("miniutils.pragma.{name} requires 'astor' to be installed to obtain source code"
-                                  .format(name=lift.__name__))
-            except Exception as ex:  # pragma: nocover
-                raise RuntimeError(astor.dump_tree(f_mod)) from ex
-        else:
-            source = None
-
-        if self.return_source:
-            return source
-        else:
-            f_mod = ast.fix_missing_locations(f_mod)
-            if self.save_source:
-                temp = tempfile.NamedTemporaryFile('w', delete=True)
-                f_file = temp.name
-            no_globals = {}
-            exec(compile(f_mod, f_file, 'exec'), no_globals)
-            func = no_globals[f_mod.body[0].name]
-            if self.save_source:
-                func.__tempfile__ = temp
-                temp.write(source)
-                temp.flush()
-            return func
+        return save_or_return_source(f_file, f_mod, {}, self.return_source, self.save_source)

pragma/unroll.py

@@ -184,7 +184,8 @@ def resolve_attr_of_slice(attr):
             resolve_attr_of_slice('upper')
             resolve_attr_of_slice('step')
         else:
-            raise TypeError(type(target.slice))
+            target.slice = self.visit(target.slice)  # The index could be anything in 3.8+
+            # raise TypeError(type(target.slice))
 
     def visit_Assign(self, node):
         for it, target in enumerate(node.targets):

pragma/utils.py

@@ -0,0 +1,34 @@
+import ast
+import tempfile
+
+import astor
+
+
+def save_or_return_source(f_file, f_mod, glbls, return_source, save_source):
+    if return_source or save_source:
+        try:
+            source = astor.to_source(f_mod)
+        except Exception as ex:  # pragma: nocover
+            raise RuntimeError(astor.dump_tree(f_mod)) from ex
+    else:
+        source = None
+
+    if return_source:
+        return source
+
+    f_mod = ast.fix_missing_locations(f_mod)
+    if save_source:
+        temp = tempfile.NamedTemporaryFile('w', delete=False)
+        f_file = temp.name
+    exec(compile(f_mod, f_file, 'exec'), glbls)
+    func = glbls[f_mod.body[0].name]
+    if save_source:
+        func.__tempfile__ = temp
+        # When there are other decorators, the co_firstlineno of *some* python distributions gets confused
+        # and thinks they will be there even when they are not written to the file, causing readline overflow
+        # So we put some empty lines to make them align
+        temp.write(source)
+        temp.write('\n' * func.__code__.co_firstlineno)
+        temp.flush()
+        temp.close()
+    return func

requirements.txt

@@ -5,3 +5,4 @@ coveralls
 coverage
 nose
 ipython
+bump2version

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='pragma',
-    version='0.2.4',
+    version='0.2.5',
     packages=['pragma', 'pragma.core', 'pragma.core.resolve'],
     url='https://github.com/scnerd/pypragma',
     license='MIT',

tests/test_collapse_literals.py

@@ -56,6 +56,7 @@ def test_repeated_decoration(self):
         @pragma.collapse_literals
         def f():
             return 2
+
         f = pragma.collapse_literals(f)
 
         result = '''
@@ -165,14 +166,15 @@ def test_side_effects_0cause(self):
         # This will never fail, but it causes other tests to fail
         # if it incorrectly moves 'a' from the closure to the module globals
         a = 0
+
         @pragma.collapse_literals
         def f():
             x = a
 
     def test_side_effects_1effect(self):
         @pragma.collapse_literals
         def f2():
-            for a in range(3): # failure occurs when this is interpreted as "for 0 in range(3)"
+            for a in range(3):  # failure occurs when this is interpreted as "for 0 in range(3)"
                 x = a
 
     def test_iteration_variable(self):
@@ -183,6 +185,7 @@ def test_iteration_variable(self):
         @pragma.collapse_literals
         def f1():
             x = glbvar
+
         result = '''
         def f1():
             x = 0
@@ -195,6 +198,7 @@ def f1():
         def f2():
             for glbvar in range(3):
                 x = glbvar
+
         result = '''
         def f2():
             for glbvar in range(3):
@@ -503,10 +507,12 @@ def f():
         self.assertEqual(f(), 4)
 
     def test_assignment_slice(self):
-        i=2
+        i = 2
+
         @pragma.collapse_literals
         def f1(x):
             x[i] = 1
+
         result = '''
         def f1(x):
             x[2] = 1
@@ -517,6 +523,7 @@ def f1(x):
         def f2(x):
             j = 1
             x[j] += 1
+
         result = '''
         def f2(x):
             j = 1
@@ -532,12 +539,14 @@ def f2(x):
 
     def test_collapse_slice_with_assign(self):
         a = 1
+
         @pragma.collapse_literals
         def f():
             x = object()
             x[a:4] = 0
             x = 2
             x[x] = 0
+
         result = '''
         def f():
             x = object()
@@ -552,6 +561,7 @@ def f():
             x = [1, 2, 0]
             x[x[x[0]]] = 3  # transformer loses certainty in literal value of x
             x[x[x[0]]] = 4  # so it is not collapsed here, but this is a nonsensical use case after all
+
         result = '''
         def f():
             x = [1, 2, 0]
@@ -563,6 +573,7 @@ def f():
 
     def test_slice_assign_(self):
         a = [1]
+
         @pragma.collapse_literals
         def f():
             x[a[0]] = 0
@@ -578,10 +589,12 @@ def f():
     def test_explicit_collapse(self):
         a = 2
         b = 3
+
         @pragma.collapse_literals(explicit_only=True, b=b)
         def f():
             x = a
             y = b
+
         result = '''
         def f():
             x = a
@@ -592,6 +605,7 @@ def f():
         @pragma.collapse_literals(explicit_only=True)
         def f():
             x = a
+
         result = '''
         def f():
             x = a
@@ -627,7 +641,7 @@ def test_mathematical_deduction(self):
         def f(x):
             yield (x / 1) + 0
             yield 0 - x
-            yield 0 * (x ** 2 + 3*x - 2)
+            yield 0 * (x ** 2 + 3 * x - 2)
             yield 0 % x
 
         result = '''

tests/test_regression_issue21.py

@@ -0,0 +1,64 @@
+import pragma
+from tests.test_pragma import PragmaTest
+
+
+class TestIssue21(PragmaTest):
+    def test_unroll_setitem(self):
+        # This is expected to work properly
+        biglist = list(range(20))
+
+        @pragma.unroll
+        def hey():
+            for i in range(10):
+                biglist.__setitem__(i, 0)
+
+        result = '''
+        def hey():
+            biglist.__setitem__(0, 0)
+            biglist.__setitem__(1, 0)
+            biglist.__setitem__(2, 0)
+            biglist.__setitem__(3, 0)
+            biglist.__setitem__(4, 0)
+            biglist.__setitem__(5, 0)
+            biglist.__setitem__(6, 0)
+            biglist.__setitem__(7, 0)
+            biglist.__setitem__(8, 0)
+            biglist.__setitem__(9, 0)
+        '''
+
+        self.assertSourceEqual(hey, result)
+        hey()
+        self.assertListEqual(
+            biglist,
+            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0] + list(range(10, 20))
+        )
+
+    def test_unroll_set_index(self):
+        # This did not work as expected
+        biglist = list(range(20))
+
+        @pragma.unroll
+        def hey2():
+            for i in range(10):
+                biglist[i] = 1
+
+        result = '''
+        def hey2():
+            biglist[0] = 1
+            biglist[1] = 1
+            biglist[2] = 1
+            biglist[3] = 1
+            biglist[4] = 1
+            biglist[5] = 1
+            biglist[6] = 1
+            biglist[7] = 1
+            biglist[8] = 1
+            biglist[9] = 1
+        '''
+
+        self.assertSourceEqual(hey2, result)
+        hey2()
+        self.assertListEqual(
+            biglist,
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1] + list(range(10, 20))
+        )