Address more of Erik's feedback

Author: lwaern-intel

py/dml/types_test.py

@@ -7,6 +7,7 @@
 from dml import expr
 from dml import traits
 from dml import logging
+from contextlib import contextmanager
 
 import unittest
 
@@ -104,29 +105,33 @@ def bitfields():
         self.assert_eq(TInt(8, False, members={}), TInt(8, False, members={}))
         self.assert_eq(bitfields(), bitfields())
 
-        t = bitfields()
-
+        @contextmanager
         def neq_bitfields_testcase():
-            nonlocal t
-            self.assert_neq(bitfields(), t)
             t = bitfields()
+            yield t
+            self.assert_neq(bitfields(), t)
 
         # Presence of members matter
-        del t.members['a']
-        neq_bitfields_testcase()
+        with neq_bitfields_testcase() as t:
+            del t.members['a']
+
         # type of members matters
-        t.members['a'][0].signed = True
-        neq_bitfields_testcase()
+        with neq_bitfields_testcase() as t:
+            t.members['a'][0].signed = True
+
         # msb/lsb of members matters
-        t.members['a'] = (TInt(8, False), 12, 5)
-        neq_bitfields_testcase()
+        with neq_bitfields_testcase() as t:
+            (mt, msb, lsb) = t.members['a']
+            t.members['a'] = (mt, msb - 1, lsb - 1)
+
         # order matters
-        t.members = dict(reversed(t.members.items()))
-        neq_bitfields_testcase()
+        with neq_bitfields_testcase() as t:
+            t.members = dict(reversed(t.members.items()))
+
         # names matter
-        t.members['c'] = t.members['b']
-        del t.members['b']
-        neq_bitfields_testcase()
+        with neq_bitfields_testcase() as t:
+            t.members['c'] = t.members['b']
+            del t.members['b']
 
     def test_TEndianInt(self):
         self.assert_eq(TEndianInt(8, False, 'big-endian'),
@@ -158,37 +163,34 @@ def test_TArray(self):
         b.size = a.size
         self.assert_eq(a, b)
 
-        a = self.mkTArray(4)
-        b = self.mkTArray(4)
-
+        @contextmanager
         def array_testcase(res):
-            nonlocal a, b
+            a = self.mkTArray(4)
+            b = self.mkTArray(4)
+            yield (a, b)
             if res:
                 self.assert_eq(a, b)
             else:
                 self.assert_neq(a, b)
 
-            a = self.mkTArray(4)
-            b = self.mkTArray(4)
-
         # Base type matters
-        a.base.signed = True
-        array_testcase(False)
+        with array_testcase(False) as (a, _):
+            a.base.signed = True
 
         # Constness: constness of base type and array are normalized
-        a.const = True
-        array_testcase(False)
-        a.const = True
-        b.const = True
-        array_testcase(True)
-        a.base.const = True
-        b.const = True
-        array_testcase(True)
-        a.base = self.mkTArray(4)
-        b.base = self.mkTArray(4)
-        a.base.base.const = True
-        b.const = True
-        array_testcase(True)
+        with array_testcase(False) as (a, _):
+            a.const = True
+        with array_testcase(True) as (a, b):
+            a.const = True
+            b.const = True
+        with array_testcase(True) as (a, b):
+            a.base.const = True
+            b.const = True
+        with array_testcase(True) as (a, b):
+            a.base = self.mkTArray(4)
+            b.base = self.mkTArray(4)
+            a.base.base.const = True
+            b.const = True
 
         # Pointers are not equivalent to arrays
         self.assert_neq(self.mkTArray(4), TPtr(TInt(32, False)))