Evaluate first comprehension iterator in outer scope (#1130)

* Evaluate first comprehension iterator in outer scope

* Add tests

* Update rewriting code path

* Add dome more tests

* Revert "Workaround for #809 (#810)"

This reverts commit 14f34c9.

* Guard against invalid comprehension iterators

Author: BCSharp

Src/IronPython/Compiler/Ast/Comprehension.cs

@@ -75,6 +75,9 @@ public sealed class ListComprehension : Comprehension {
         private readonly ComprehensionIterator[] _iterators;
 
         public ListComprehension(Expression item, ComprehensionIterator[] iterators) {
+            if (iterators is null || iterators.Length < 1) throw new ArgumentException("comprehension with no generators");
+            if (iterators[0] is not ComprehensionFor) throw new ArgumentException("comprehension with invalid generator");
+
             Item = item;
             _iterators = iterators;
             Scope = new ComprehensionScope(this);
@@ -121,6 +124,9 @@ public sealed class SetComprehension : Comprehension {
         private readonly ComprehensionIterator[] _iterators;
 
         public SetComprehension(Expression item, ComprehensionIterator[] iterators) {
+            if (iterators is null || iterators.Length < 1) throw new ArgumentException("comprehension with no generators");
+            if (iterators[0] is not ComprehensionFor) throw new ArgumentException("comprehension with invalid generator");
+
             Item = item;
             _iterators = iterators;
             Scope = new ComprehensionScope(this);
@@ -167,6 +173,9 @@ public sealed class DictionaryComprehension : Comprehension {
         private readonly ComprehensionIterator[] _iterators;
 
         public DictionaryComprehension(Expression key, Expression value, ComprehensionIterator[] iterators) {
+            if (iterators is null || iterators.Length < 1) throw new ArgumentException("comprehension with no generators");
+            if (iterators[0] is not ComprehensionFor) throw new ArgumentException("comprehension with invalid generator");
+
             Key = key;
             Value = value;
             _iterators = iterators;

Src/IronPython/Compiler/Ast/PythonAst.cs

@@ -720,6 +720,7 @@ internal PythonAst MakeLookupCode() {
 
         internal class LookupVisitor : MSAst.ExpressionVisitor {
             private readonly MSAst.Expression _globalContext;
+            private readonly Dictionary<MSAst.Expression, ScopeStatement> _outerComprehensionScopes = new();
             private ScopeStatement _curScope;
 
             public LookupVisitor(PythonAst ast, MSAst.Expression globalContext) {
@@ -739,6 +740,12 @@ protected override MSAst.Expression VisitExtension(MSAst.Expression node) {
                     return PythonAst._globalContext;
                 }
 
+                // outer comprehension iterable is visited in outer comprehension scope
+                if (_outerComprehensionScopes.TryGetValue(node, out ScopeStatement outerComprehensionScope)) {
+                    _outerComprehensionScopes.Remove(node);
+                    return VisitExtensionInScope(node, outerComprehensionScope);
+                }
+
                 // we need to re-write nested scopes
                 if (node is ScopeStatement scope) {
                     return base.VisitExtension(VisitScope(scope));
@@ -815,14 +822,27 @@ private MSAst.Expression VisitComprehension(Comprehension comprehension) {
 
                 ScopeStatement prevScope = _curScope;
                 try {
-                    // rewrite the comprehension in a new scope
+                    // mark the first (outermost) "for" iterator for rewrite in the current scope
+                    _outerComprehensionScopes[((ComprehensionFor)comprehension.Iterators[0]).List] = _curScope;
+
+                    // rewrite the rest of comprehension in the new scope
                     _curScope = newScope;
 
                     return base.VisitExtension(newComprehension);
                 } finally {
                     _curScope = prevScope;
                 }
             }
+
+            private MSAst.Expression VisitExtensionInScope(MSAst.Expression node, ScopeStatement scope) {
+                ScopeStatement prevScope = _curScope;
+                try {
+                    _curScope = scope;
+                    return VisitExtension(node); // not base.VisitExtension
+                } finally {
+                    _curScope = prevScope;
+                }
+            }
         }
 
         #endregion

Src/IronPython/Compiler/Ast/PythonNameBinder.cs

@@ -216,7 +216,7 @@ public override bool Walk(ClassDefinition node) {
                     dec.Walk(this);
                 }
             }
-            
+
             PushScope(node);
 
             node.ModuleNameVariable = _globalScope.EnsureGlobalVariable("__name__");
@@ -248,10 +248,30 @@ public override bool Walk(DelStatement node) {
             return true;
         }
 
-        public override bool Walk(ListComprehension node) {
+        // Comprehensions
+
+        private void WalkComprehensionIterators(Comprehension node) {
             node.Parent = _currentScope;
+
+            // Special walk case: first (outermost) "for" iterator
+            // See also: PythonAst.LookupVisitor.VisitComprehension(...)
+            var outermostFor = (ComprehensionFor)node.Iterators[0];
+            outermostFor.List.Walk(this);
             PushScope(node.Scope);
-            return base.Walk(node);
+            Walk(outermostFor);
+            outermostFor.Left.Walk(this);
+            PostWalk(outermostFor);
+
+            // Regular walk cases: remaining iterators/conditionals
+            for (int i = 1; i < node.Iterators.Count; i++) {
+                node.Iterators[i].Walk(this);
+            }
+        }
+
+        public override bool Walk(ListComprehension node) {
+            WalkComprehensionIterators(node);
+            node.Item.Walk(this);
+            return false;
         }
 
         public override void PostWalk(ListComprehension node) {
@@ -264,9 +284,9 @@ public override void PostWalk(ListComprehension node) {
         }
 
         public override bool Walk(SetComprehension node) {
-            node.Parent = _currentScope;
-            PushScope(node.Scope);
-            return base.Walk(node);
+            WalkComprehensionIterators(node);
+            node.Item.Walk(this);
+            return false;
         }
 
         public override void PostWalk(SetComprehension node) {
@@ -279,9 +299,10 @@ public override void PostWalk(SetComprehension node) {
         }
 
         public override bool Walk(DictionaryComprehension node) {
-            node.Parent = _currentScope;
-            PushScope(node.Scope);
-            return base.Walk(node);
+            WalkComprehensionIterators(node);
+            node.Key.Walk(this);
+            node.Value.Walk(this);
+            return false;
         }
 
         public override void PostWalk(DictionaryComprehension node) {

Src/StdLib/Lib/pickle.py

@@ -177,10 +177,8 @@ def __init__(self, value):
 MEMOIZE          = b'\x94'  # store top of the stack in memo
 FRAME            = b'\x95'  # indicate the beginning of a new frame
 
-# _dir is a workaround for https://github.com/IronLanguages/ironpython3/issues/809
-_dir = dir()
-__all__.extend([x for x in _dir if re.match("[A-Z][A-Z0-9_]+$", x)])
-del _dir
+__all__.extend([x for x in dir() if re.match("[A-Z][A-Z0-9_]+$", x)])
+
 
 class _Framer:
 

Tests/test_dictcomp.py

@@ -99,4 +99,27 @@ def eval(f, i):
         r = {k:eval(lambda i: i+v, k+v) for k in range(2)}
         self.assertEqual(r, {0:4, 1:5})
 
+    def test_ipy3_gh809(self):
+        """https://github.com/IronLanguages/ironpython3/issues/809"""
+
+        # iterable is evaluated in the outer scope
+        self.assertIn('self', {x : None for x in dir()})
+
+        # this rule applies recursively to nested comprehensions
+        self.assertIn('self', {x : None for x in {y : None for y in dir()}})
+        self.assertIn('self', {x : None for x in {y : None for y in {z : None for z in dir()}}})
+
+        # this only applies to the first iterable
+        # subsequent iterables are evaluated within the comprehension scope
+        self.assertEqual({(0, 'x') : None}, {(x, y) : None for x in range(1) for y in dir() if not y.startswith('.')}) # (filtering out auxiliary variable staring with a dot, used by CPython)
+
+        # also subsequent conditions are evaluated within the comprehension scope
+        a, b, c, d = range(4)
+        self.assertTrue(len(dir()) >= 4)
+        self.assertEqual({}, {None : dir() for x in range(1) if len(dir()) >= 4})
+
+        # lambdas create a new scope
+        self.assertEqual({}, {x : None for x in (lambda: dir())()})
+        self.assertEqual({None: []}, {None: x() for x in {lambda: dir() : None}})
+
 run_test(__name__)

Tests/test_listcomp.py

@@ -4,7 +4,7 @@
 
 import unittest
 
-from iptest import run_test
+from iptest import run_test, is_cli
 
 class ListCompTest(unittest.TestCase):
     def test_positive(self):
@@ -41,4 +41,69 @@ def test_negative(self):
         self.assertRaises(NameError, lambda: [(x, z) for x in "iron" if z > x for z in "python" ])
         self.assertRaises(NameError, lambda: [(i, j) for i in range(10) if j < 'c' for j in ['a', 'b', 'c'] if i % 3 == 0])
 
+    def test_ipy3_gh809(self):
+        """https://github.com/IronLanguages/ironpython3/issues/809"""
+
+        # iterable is evaluated in the outer scope
+        self.assertIn('self', [x for x in dir()])
+
+        # this rule applies recursively to nested comprehensions
+        self.assertIn('self', [x for x in [y for y in dir()]])
+        self.assertIn('self', [x for x in [y for y in [z for z in dir()]]])
+
+        # this only applies to the first iterable
+        # subsequent iterables are evaluated within the comprehension scope
+        self.assertEqual([(0, 'x')], [(x, y) for x in range(1) for y in dir() if not y.startswith('.')]) # (filtering out auxiliary variable staring with a dot, used by CPython)
+
+        # also subsequent conditions are evaluated within the comprehension scope
+        a, b, c, d = range(4)
+        self.assertTrue(len(dir()) >= 4)
+        self.assertEqual([], [dir() for x in range(1) if len(dir()) >= 4])
+
+        # subsequent iterables introduce local variables after the first iteration
+        self.assertEqual([(0, 'x'), (1, 'x'), (1, 'y'), (2, 'x'), (2, 'y')],
+                         [(x, y) for x in range(3)
+                                 for y in dir() if not y.startswith('.')])
+        self.assertEqual([(0, 'x', 'x'), (0, 'x', 'y'), 
+                          (1, 'x', 'x'), (1, 'x', 'y'), (1, 'x', 'z'), 
+                          (1, 'y', 'x'), (1, 'y', 'y'), (1, 'y', 'z'), 
+                          (1, 'z', 'x'), (1, 'z', 'y'), (1, 'z', 'z')],
+                         [(x, y, z) for x in range(2)
+                                    for y in dir() if not y.startswith('.')
+                                    for z in dir() if not z.startswith('.')])
+
+        # lambdas create a new scope
+        self.assertEqual([], [x for x in (lambda: dir())()])
+        self.assertEqual([[]], [x() for x in [lambda: dir()]])
+
+        # first iterable is captured and subsequent assignments do not change it
+        self.maxDiff = None
+        x, y, z = range(1), range(2), range(3)
+        if is_cli:
+            _dir = ['x', 'y']
+            # TODO: should be: _dir = ['x', 'y', 'z']
+            # See: https://github.com/IronLanguages/ironpython3/issues/1132
+        else:
+            _dir = ['.0', 'x', 'y', 'z'] # adds implementation-level variable '.0'
+        # below, x and first/third y is local, second y and z is from the outer scope
+        self.assertEqual([(x, y, z, dir()) for x in y for y in z],
+                         [(0, 0, range(3), _dir),
+                          (0, 1, range(3), _dir),
+                          (0, 2, range(3), _dir),
+                          (1, 0, range(3), _dir),
+                          (1, 1, range(3), _dir),
+                          (1, 2, range(3), _dir)])
+
+        # mixing scopes
+        def apply(f, i): return f(i)
+        x = 2
+        res = [x for x in apply(lambda i: range(i+x), x)]
+        self.assertEqual(res, [0, 1, 2, 3])
+
+        res = [(x, y) for x in apply(lambda i: range(i+x), x) for y in apply(lambda i: range(i+x//2), x)]
+        self.assertEqual(res, [(1, 0), (2, 0), (2, 1), (2, 2), (3, 0), (3, 1), (3, 2), (3, 3)])
+
+        res = [x for x in [y for y in apply(lambda i: range(i+x), x)]]
+        self.assertEqual(res, [0, 1, 2, 3])
+
 run_test(__name__)

Tests/test_setcomp.py

@@ -99,4 +99,26 @@ def eval(f, i):
         r = {eval(lambda i: i+v, k+v) for k in range(2)}
         self.assertEqual(r, set([4, 5]))
 
+    def test_ipy3_gh809(self):
+        """https://github.com/IronLanguages/ironpython3/issues/809"""
+
+        # iterable is evaluated in the outer scope
+        self.assertIn('self', {x for x in dir()})
+
+        # this rule applies recursively to nested comprehensions
+        self.assertIn('self', {x for x in {y for y in dir()}})
+        self.assertIn('self', {x for x in {y for y in {z for z in dir()}}})
+
+        # this only applies to the first iterable
+        # subsequent iterables are evaluated within the comprehension scope
+        self.assertEqual({(0, 'x')}, {(x, y) for x in range(1) for y in dir() if not y.startswith('.')}) # (filtering out auxiliary variable staring with a dot, used by CPython)
+
+        # also subsequent conditions are evaluated within the comprehension scope
+        a, b, c, d = range(4)
+        self.assertTrue(len(dir()) >= 4)
+        self.assertEqual(set(), {dir() for x in range(1) if len(dir()) >= 4})
+
+        # lambdas create a new scope
+        self.assertEqual(set(), {x for x in (lambda: dir())()})
+
 run_test(__name__)