From 5cebe752e8b668a9679a18f5c8680df5d569084f Mon Sep 17 00:00:00 2001
From: David Maxson <David.Maxson@jhuapl.edu>
Date: Mon, 3 May 2021 07:20:54 -0700
Subject: [PATCH 1/5] Fixed bug actually with lift that was breaking tests

---
 pragma/core/transformer.py       | 31 +---------------
 pragma/lift.py                   | 29 +--------------
 pragma/utils.py                  | 34 +++++++++++++++++
 tests/test_regression_issue21.py | 64 ++++++++++++++++++++++++++++++++
 4 files changed, 102 insertions(+), 56 deletions(-)
 create mode 100644 pragma/utils.py
 create mode 100644 tests/test_regression_issue21.py

diff --git a/pragma/core/transformer.py b/pragma/core/transformer.py
index 7ea9019..2e07055 100644
--- a/pragma/core/transformer.py
+++ b/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
 
diff --git a/pragma/lift.py b/pragma/lift.py
index ac99e89..c7ffc52 100644
--- a/pragma/lift.py
+++ b/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)
diff --git a/pragma/utils.py b/pragma/utils.py
new file mode 100644
index 0000000..fa86531
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/tests/test_regression_issue21.py b/tests/test_regression_issue21.py
new file mode 100644
index 0000000..68e1154
--- /dev/null
+++ b/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))
+        )

From 1f25c192f184344d0e29d9e23c11d107a21bac7f Mon Sep 17 00:00:00 2001
From: David Maxson <David.Maxson@jhuapl.edu>
Date: Mon, 3 May 2021 07:25:21 -0700
Subject: [PATCH 2/5] Added bump2version support

---
 .bumpversion.cfg | 7 +++++++
 requirements.txt | 1 +
 2 files changed, 8 insertions(+)
 create mode 100644 .bumpversion.cfg

diff --git a/.bumpversion.cfg b/.bumpversion.cfg
new file mode 100644
index 0000000..1888538
--- /dev/null
+++ b/.bumpversion.cfg
@@ -0,0 +1,7 @@
+[bumpversion]
+current_version = 0.2.4
+commit = True
+tag = True
+
+[bumpversion:file:setup.py]
+[bumpversion:file:docs/conf.py]
diff --git a/requirements.txt b/requirements.txt
index 71baec6..dc07c60 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,3 +5,4 @@ coveralls
 coverage
 nose
 ipython
+bump2version

From dbc63fa29525cf5b08307656483c3d8149cedacd Mon Sep 17 00:00:00 2001
From: David Maxson <David.Maxson@jhuapl.edu>
Date: Mon, 3 May 2021 07:59:59 -0700
Subject: [PATCH 3/5] Fixed 3.8/3.9 compatibility issue with .Index being
 replaced by the underlying ast types

---
 docs/todo.rst                   |  1 +
 pragma/collapse_literals.py     |  3 ++-
 pragma/unroll.py                |  3 ++-
 tests/test_collapse_literals.py | 20 +++++++++++++++++---
 4 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/docs/todo.rst b/docs/todo.rst
index 1d49e58..5f2edcc 100644
--- a/docs/todo.rst
+++ b/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::
diff --git a/pragma/collapse_literals.py b/pragma/collapse_literals.py
index 48e2869..96fa56a 100644
--- a/pragma/collapse_literals.py
+++ b/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):
diff --git a/pragma/unroll.py b/pragma/unroll.py
index 468aee9..adeb229 100644
--- a/pragma/unroll.py
+++ b/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):
diff --git a/tests/test_collapse_literals.py b/tests/test_collapse_literals.py
index 3a648f2..554ae76 100644
--- a/tests/test_collapse_literals.py
+++ b/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,6 +166,7 @@ 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
@@ -172,7 +174,7 @@ def f():
     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 = '''

From dd413339ace025d255d75b5bd8ec47720674667d Mon Sep 17 00:00:00 2001
From: David Maxson <David.Maxson@jhuapl.edu>
Date: Mon, 3 May 2021 08:00:17 -0700
Subject: [PATCH 4/5] =?UTF-8?q?Bump=20version:=200.2.4=20=E2=86=92=200.2.5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .bumpversion.cfg | 3 ++-
 docs/conf.py     | 2 +-
 setup.py         | 2 +-
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/.bumpversion.cfg b/.bumpversion.cfg
index 1888538..392f7db 100644
--- a/.bumpversion.cfg
+++ b/.bumpversion.cfg
@@ -1,7 +1,8 @@
 [bumpversion]
-current_version = 0.2.4
+current_version = 0.2.5
 commit = True
 tag = True
 
 [bumpversion:file:setup.py]
+
 [bumpversion:file:docs/conf.py]
diff --git a/docs/conf.py b/docs/conf.py
index a7cf34f..c44fdcc 100644
--- a/docs/conf.py
+++ b/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.
diff --git a/setup.py b/setup.py
index 35fa9d8..ed039cd 100644
--- a/setup.py
+++ b/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',

From 33824c49df0ca23ffa53b20fd4ab005370101331 Mon Sep 17 00:00:00 2001
From: David Maxson <David.Maxson@jhuapl.edu>
Date: Mon, 3 May 2021 08:04:13 -0700
Subject: [PATCH 5/5] Add 3.9 to the test suite

---
 .travis.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.travis.yml b/.travis.yml
index 7ac0639..ab6644c 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -4,6 +4,7 @@ python:
  - "3.6"
  - "3.7"
  - "3.8"
+ - "3.9"
 dist: xenial
 
 install: