more tests

Author: iritkatriel

Lib/test/test_listcomps.py

@@ -530,6 +530,34 @@ def g(n):
         self.assertEqual(g.__code__.co_cellvars, ("x",))
         self.assertEqual([fn() for fn in g(2)[1]], [1, 1])
 
+    def test_nested_inlined_async_comp_iter_var_is_fast_local(self):
+        import asyncio
+
+        async def agen(n):
+            for i in range(n):
+                yield i
+
+        async def f(n):
+            return [[x async for _ in agen(2)] async for x in agen(n)]
+
+        self.assertEqual(f.__code__.co_cellvars, ())
+        self.assertEqual(asyncio.run(f(2)), [[0, 0], [1, 1]])
+
+        async def g(n):
+            return [[(lambda: x) async for _ in agen(2)] async for x in agen(n)]
+
+        self.assertEqual(g.__code__.co_cellvars, ("x",))
+        out = asyncio.run(g(2))
+        self.assertEqual([fn() for fn in out[1]], [1, 1])
+
+    def test_inlined_comprehension_name_mangling_in_method_scope(self):
+        class C:
+            def f(self):
+                __x = 42
+                return [__x for _ in (0,)]
+
+        self.assertEqual(C().f(), [42])
+
     def test_nested_references___class__(self):
         code = """
             res = [[__class__ for _ in (0,)] for _ in (1,)]
@@ -914,6 +942,46 @@ def inner():
             {"snaps": [1, 2], "vals": [2, 2], "consistent": [True, True]},
             ns={"sys": sys}, scopes=["module", "function"])
 
+    def test_frame_locals_nested_comp_cell_and_enclosing_free(self):
+        # Stress a nested inlined shape where a comp cell and enclosing free
+        # share a name; all f_locals views must stay consistent.
+        code = """
+            def outer(x):
+                def inner():
+                    return [(
+                             lambda: x,
+                             [[x for _ in (0,)] for _ in (0,)][0][0],
+                             dict(**sys._getframe().f_locals),
+                             len(sys._getframe().f_locals),
+                             list(sys._getframe().f_locals.keys()),
+                             list(sys._getframe().f_locals.values()),
+                             list(sys._getframe().f_locals.items()),
+                             dict(sys._getframe().f_locals.items()))
+                            for x in x]
+                return inner()
+            result = outer([1, 2])
+            snaps = [d['x'] for _, _, d, *_ in result]
+            vals = [fn() for fn, *_ in result]
+            nested_vals = [nested for _, nested, *_ in result]
+            consistent = []
+            for _, _, d, n, ks, vs, it, d_items in result:
+                consistent.append(
+                    n == len(ks) == len(vs) == len(it)
+                    and ks.count('x') == 1
+                    and d == d_items == dict(zip(ks, vs))
+                )
+        """
+        import sys
+        self._check_in_scopes(
+            code,
+            {
+                "snaps": [1, 2],
+                "vals": [2, 2],
+                "nested_vals": [1, 2],
+                "consistent": [True, True],
+            },
+            ns={"sys": sys}, scopes=["module", "function"])
+
     def _recursive_replace(self, maybe_code):
         if not isinstance(maybe_code, types.CodeType):
             return maybe_code

Lib/test/test_symtable.py

@@ -558,6 +558,36 @@ def test_inlined_comprehension_use_of_enclosing_free_in_function(self):
         self.assertTrue(comp.lookup("x").is_free())
         self.assertTrue(comp.lookup("x").is_referenced())
 
+    def test_inlined_comprehension_use_of_nonlocal_in_function(self):
+        st = symtable.symtable(
+            "def outer():\n"
+            "    x = 1\n"
+            "    def inner():\n"
+            "        nonlocal x\n"
+            "        return [x for _ in ()]",
+            "?", "exec")
+        inner = find_block(find_block(st, "outer"), "inner")
+        self.assertTrue(inner.lookup("x").is_nonlocal())
+        self.assertTrue(inner.lookup("x").is_free())
+        self.assertTrue(inner.lookup("x").is_referenced())
+        comp, = inner.get_children()
+        self.assertTrue(comp.lookup("x").is_free())
+        self.assertTrue(comp.lookup("x").is_referenced())
+
+    def test_inlined_comprehension_use_of_explicit_global_in_function(self):
+        st = symtable.symtable(
+            "def f():\n"
+            "    global g\n"
+            "    return [g for _ in ()]",
+            "?", "exec")
+        f = find_block(st, "f")
+        self.assertTrue(f.lookup("g").is_global())
+        self.assertTrue(f.lookup("g").is_declared_global())
+        self.assertFalse(f.lookup("g").is_referenced())
+        comp, = f.get_children()
+        self.assertTrue(comp.lookup("g").is_global())
+        self.assertTrue(comp.lookup("g").is_referenced())
+
     def test_inlined_comprehension_nested_function_use_not_on_enclosing(self):
         # The load of x is in the lambda's code object, not inner's.
         st = symtable.symtable(
@@ -661,6 +691,21 @@ def test_inlined_nested_comprehension_iter_var_is_fast_local(self):
         self.assertTrue(inner.lookup("x").is_free())
         self.assertTrue(inner.lookup("x").is_referenced())
 
+    def test_inlined_nested_mixed_comprehension_iter_var_is_referenced(self):
+        st = symtable.symtable(
+            "def f():\n"
+            "    return [{x for _ in (0,)} for x in (42,)]",
+            "?", "exec")
+        f = find_block(st, "f")
+        with self.assertRaises(KeyError):
+            f.lookup("x")
+        outer, = f.get_children()
+        self.assertTrue(outer.lookup("x").is_local())
+        self.assertTrue(outer.lookup("x").is_referenced())
+        inner, = outer.get_children()
+        self.assertTrue(inner.lookup("x").is_free())
+        self.assertTrue(inner.lookup("x").is_referenced())
+
     def test_inlined_nested_comprehension_captured_iter_is_cell(self):
         st = symtable.symtable(
             "def f():\n"