[pypy-commit] pypy py3k: Implement super() and all its magic:

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: py3k
Changeset: r47981:a8cfd510a295
Date: 2011-10-12 22:19 +0200
http://bitbucket.org/pypy/pypy/changeset/a8cfd510a295/

Log:	Implement super() and all its magic: There is a hidden __class__
	cell, and self is the first argument of the function.

diff --git a/pypy/interpreter/astcompiler/codegen.py b/pypy/interpreter/astcompiler/codegen.py
--- a/pypy/interpreter/astcompiler/codegen.py
+++ b/pypy/interpreter/astcompiler/codegen.py
@@ -201,7 +201,7 @@
                           filename=self.compile_info.filename)
 
     def name_op(self, identifier, ctx):
-        """Generate an operation appropiate for the scope of the identifier."""
+        """Generate an operation appropriate for the scope of the identifier."""
         scope = self.scope.lookup(identifier)
         op = ops.NOP
         container = self.names
@@ -1248,10 +1248,10 @@
         self._handle_body(cls.body)
         # return the (empty) __class__ cell
         scope = self.scope.lookup("@__class__")
-        if scope == symtable.SCOPE_UNKNOWN:
+        if scope == symtable.SCOPE_CELL:
+            # Return the cell where to store __class__
+            self.emit_op_arg(ops.LOAD_CLOSURE, self.cell_vars["@__class__"])
+        else:
             # This happens when nobody references the cell
             self.load_const(self.space.w_None)
-        else:
-            # Return the cell where to store __class__
-            self.emit_op_arg(ops.LOAD_CLOSURE, self.cell_vars["@__class__"])
         self.emit_op(ops.RETURN_VALUE)
diff --git a/pypy/interpreter/astcompiler/symtable.py b/pypy/interpreter/astcompiler/symtable.py
--- a/pypy/interpreter/astcompiler/symtable.py
+++ b/pypy/interpreter/astcompiler/symtable.py
@@ -163,6 +163,8 @@
             self._finalize_name(name, flags, local, bound, free, globs)
         if not self._hide_bound_from_nested_scopes:
             self._pass_on_bindings(local, bound, globs, new_bound, new_globs)
+        else:
+            self._pass_special_names(local, new_bound)
         child_frees = {}
         for child in self.children:
             # Symbol dictionaries are copied to avoid having child scopes
@@ -214,6 +216,12 @@
         self.import_star = None
         self.bare_exec = None
 
+    def note_symbol(self, identifier, role):
+        # Special-case super: it counts as a use of __class__
+        if role == SYM_USED and identifier == 'super':
+            self.note_symbol('@__class__', SYM_USED)
+        Scope.note_symbol(self, identifier, role)
+
     def note_yield(self, yield_node):
         if self.return_with_value:
             raise SyntaxError("'return' with argument inside generator",
@@ -295,6 +303,16 @@
     def mangle(self, name):
         return misc.mangle(name, self.name)
 
+    def _pass_special_names(self, local, new_bound):
+        assert '@__class__' in local
+        new_bound['@__class__'] = None
+
+    def _finalize_cells(self, free):
+        for name, role in self.symbols.iteritems():
+            if role == SCOPE_LOCAL and name in free and name == '@__class__':
+                self.symbols[name] = SCOPE_CELL
+                del free[name]
+
 
 class SymtableBuilder(ast.GenericASTVisitor):
     """Find symbol information from AST."""
@@ -373,6 +391,8 @@
         self.visit_sequence(clsdef.bases)
         self.visit_sequence(clsdef.decorator_list)
         self.push_scope(ClassScope(clsdef), clsdef)
+        self.note_symbol('@__class__', SYM_ASSIGNED)
+        self.note_symbol('__locals__', SYM_PARAM)
         self.visit_sequence(clsdef.body)
         self.pop_scope()
 
diff --git a/pypy/module/__builtin__/descriptor.py b/pypy/module/__builtin__/descriptor.py
--- a/pypy/module/__builtin__/descriptor.py
+++ b/pypy/module/__builtin__/descriptor.py
@@ -47,7 +47,36 @@
         return space.call_function(object_getattribute(space),
                                    w(self), w(name))
 
-def descr_new_super(space, w_subtype, w_starttype, w_obj_or_type=None):
+def descr_new_super(space, w_subtype, w_starttype=None, w_obj_or_type=None):
+    if space.is_w(w_starttype, space.w_None):
+        # Call super(), without args -- fill in from __class__
+        # and first local variable on the stack.
+        ec = space.getexecutioncontext()
+        frame = ec.gettopframe()
+        code = frame.pycode
+        if not code:
+            raise OperationError(space.w_SystemError, space.wrap(
+                    "super(): no code object"))
+        if code.co_argcount == 0:
+            raise OperationError(space.w_SystemError, space.wrap(
+                    "super(): no arguments"))
+        w_obj = frame.locals_stack_w[0]
+        if not w_obj:
+            raise OperationError(space.w_SystemError, space.wrap(
+                    "super(): arg[0] deleted"))
+        index = 0
+        for name in code.co_freevars:
+            if name == "@__class__":
+                break
+            index += 1
+        else:
+            raise OperationError(space.w_SystemError, space.wrap(
+                    "super(): __class__ cell not found"))
+        # a kind of LOAD_DEREF
+        cell = frame.cells[len(code.co_cellvars) + index]
+        w_starttype = cell.get()
+        w_obj_or_type = w_obj
+
     if space.is_w(w_obj_or_type, space.w_None):
         w_type = None  # unbound super object
     else:
diff --git a/pypy/module/__builtin__/test/test_descriptor.py b/pypy/module/__builtin__/test/test_descriptor.py
--- a/pypy/module/__builtin__/test/test_descriptor.py
+++ b/pypy/module/__builtin__/test/test_descriptor.py
@@ -73,6 +73,23 @@
         assert d.f() == "DBCA"
         assert D.__mro__ == (D, B, C, A, object)
 
+    def test_super_magic(self):
+        class A(object):
+            def f(self):
+                return 'A'
+        class B(A):
+            def f(self):
+                return 'B' + super().f()
+        class C(A):
+            def f(self):
+                return 'C' + super().f()
+        class D(B, C):
+            def f(self):
+                return 'D' + super().f()
+        d = D()
+        assert d.f() == "DBCA"
+        assert D.__mro__ == (D, B, C, A, object)
+
     def test_super_metaclass(self):
         class xtype(type):
             def __init__(self, name, bases, dict):