Get super() and __class__ both works in non-global context. (#721)

Author: in-code-i-trust

Src/IronPython/Compiler/Ast/PythonNameBinder.cs

@@ -821,10 +821,28 @@ public override bool Walk(CallExpression node) {
             node.Parent = _currentScope;
 
             if (node.Target is NameExpression nameExpr && nameExpr.Name == "super" && _currentScope is FunctionDefinition func) {
-                _currentScope.Reference("__class__");
                 if (node.Args.Length == 0 && func.ParameterNames.Length > 0) {
-                    node.ImplicitArgs.Add(new Arg(new NameExpression("__class__")));
-                    node.ImplicitArgs.Add(new Arg(new NameExpression(func.ParameterNames[0])));
+                    if (ShouldExpandSuperSyntaxSugar(node)) {
+                        // if `super()` is referenced in a class method.
+                        _currentScope.Reference(node.Parent.Parent.Name);
+                        node.ImplicitArgs.Add(new Arg(new NameExpression(node.Parent.Parent.Name)));
+                        node.ImplicitArgs.Add(new Arg(new NameExpression(func.ParameterNames[0])));
+                    } else {
+                        // otherwise, fallback to default implementation.
+                        _currentScope.Reference("__class__");
+                        node.ImplicitArgs.Add(new Arg(new NameExpression("__class__")));
+                        node.ImplicitArgs.Add(new Arg(new NameExpression(func.ParameterNames[0])));
+                    }
+                }
+
+                bool ShouldExpandSuperSyntaxSugar(CallExpression node) {
+                    if (!(node.Parent is FunctionDefinition)) {
+                        return false;
+                    }
+                    if (!(node.Parent.Parent is ClassDefinition)) {
+                        return false;
+                    }
+                    return true;
                 }
             }
             return base.Walk(node);

Src/IronPython/Compiler/Parser.cs

@@ -1962,8 +1962,28 @@ private Expression ParseAtom() {
                     NextToken();
                     string name = (string)t.Value;
                     _sink?.StartName(GetSourceSpan(), name);
-                    ret = new NameExpression(FixName(name));
+                    var fixedName = FixName(name);
+                    if (ShouldReplaceClassKeyword(fixedName)) {
+                        // if `__class__` variable is used in a class method, replace with `name_of_the_class`.
+                        ret = new NameExpression(CurrentClass.Name);
+                    } else {
+                        ret = new NameExpression(fixedName);
+                    }
                     ret.SetLoc(_globalParent, GetStart(), GetEnd());
+
+                    bool ShouldReplaceClassKeyword(string fixedName) {
+                        if (CurrentClass == null || CurrentFunction == null) {
+                            return false;
+                        }
+                        if (PeekToken(TokenKind.Assign)) {
+                            // avoid replacing statement like `__class__ = getproperty(...)`.
+                            return false;
+                        }
+                        if (fixedName != "__class__") {
+                            return false;
+                        }
+                        return true;
+                    }
                     return ret;
                 case TokenKind.Constant:        // literal
                     NextToken();

Tests/test_dict.py

@@ -55,12 +55,10 @@ def __setitem__(self, *args):
                 super(MyDict, self).__setitem__(*args)
 
         a = MyDict()
-        with self.assertRaises(SystemError): # TODO: remove assertRaises when https://github.com/IronLanguages/ironpython3/issues/456 is fixed
-            a[0] = 'abc'
-            self.assertEqual(a[0], 'abc')
-        with self.assertRaises(SystemError): # TODO: remove assertRaises when https://github.com/IronLanguages/ironpython3/issues/456 is fixed
-            a[None] = 3
-            self.assertEqual(a[None], 3)
+        a[0] = 'abc'
+        self.assertEqual(a[0], 'abc')
+        a[None] = 3
+        self.assertEqual(a[None], 3)
 
         class MyDict(dict):
             def __setitem__(self, *args):

Tests/test_regressions.py

@@ -996,16 +996,7 @@ def __init__(self):
             super(two, self).__init__()
             self.cnt += 1
 
-        if is_cli:
-            try:
-                self.assertEqual(two().cnt, 3)
-            except SystemError:
-                # https://github.com/IronLanguages/ironpython3/issues/451
-                pass
-            else:
-                self.fail("delete the try/except when https://github.com/IronLanguages/ironpython3/issues/451 is fixed")
-        else:
-            self.assertEqual(two().cnt, 3)
+        self.assertEqual(two().cnt, 3)
 
     def test_ipy2_gh292(self):
         """https://github.com/IronLanguages/ironpython2/issues/292"""
@@ -1449,4 +1440,32 @@ def run(self):
             self.assertEqual(thread.thread_executed, 1)
         test_thread()
 
+    def test_ipy3_gh451(self):
+        """https://github.com/IronLanguages/ironpython3/issues/451"""
+        def test():
+            class two(object):
+                def __init__(self):
+                    super().__init__()
+                    self.cnt = 1
+            return two().cnt
+        self.assertEqual(test(), 1)
+
+        class test__class__keyword(object):
+            def __new__(cls):
+                return super().__new__(cls)
+            def __init__(self):
+                super().__init__()
+            def get_self_class(self):
+                return self.__class__
+            def get_class(self):
+                return __class__
+            @classmethod
+            def get_class_class(cls):
+                return cls
+
+        o = test__class__keyword()
+        self.assertEqual(o.get_self_class(), o.get_class())
+        self.assertEqual(o.get_class(), o.get_class_class())
+
+
 run_test(__name__)