bpo-32226: Implementation of PEP 560 (core components) (#4732)

This part of the PEP implementation adds support for
__mro_entries__ and __class_getitem__ by updating
__build_class__ and PyObject_GetItem.

7 files changed, 492 insertions, 5 deletions

diff --git a/Lib/test/test_genericclass.py b/Lib/test/test_genericclass.py
new file mode 100644
index 0000000..214527b
--- /dev/null
+++ b/Lib/test/test_genericclass.py
@@ -0,0 +1,252 @@
+import unittest
+
+
+class TestMROEntry(unittest.TestCase):
+    def test_mro_entry_signature(self):
+        tested = []
+        class B: ...
+        class C:
+            def __mro_entries__(self, *args, **kwargs):
+                tested.extend([args, kwargs])
+                return (C,)
+        c = C()
+        self.assertEqual(tested, [])
+        class D(B, c): ...
+        self.assertEqual(tested[0], ((B, c),))
+        self.assertEqual(tested[1], {})
+
+    def test_mro_entry(self):
+        tested = []
+        class A: ...
+        class B: ...
+        class C:
+            def __mro_entries__(self, bases):
+                tested.append(bases)
+                return (self.__class__,)
+        c = C()
+        self.assertEqual(tested, [])
+        class D(A, c, B): ...
+        self.assertEqual(tested[-1], (A, c, B))
+        self.assertEqual(D.__bases__, (A, C, B))
+        self.assertEqual(D.__orig_bases__, (A, c, B))
+        self.assertEqual(D.__mro__, (D, A, C, B, object))
+        d = D()
+        class E(d): ...
+        self.assertEqual(tested[-1], (d,))
+        self.assertEqual(E.__bases__, (D,))
+
+    def test_mro_entry_none(self):
+        tested = []
+        class A: ...
+        class B: ...
+        class C:
+            def __mro_entries__(self, bases):
+                tested.append(bases)
+                return ()
+        c = C()
+        self.assertEqual(tested, [])
+        class D(A, c, B): ...
+        self.assertEqual(tested[-1], (A, c, B))
+        self.assertEqual(D.__bases__, (A, B))
+        self.assertEqual(D.__orig_bases__, (A, c, B))
+        self.assertEqual(D.__mro__, (D, A, B, object))
+        class E(c): ...
+        self.assertEqual(tested[-1], (c,))
+        self.assertEqual(E.__bases__, (object,))
+        self.assertEqual(E.__orig_bases__, (c,))
+        self.assertEqual(E.__mro__, (E, object))
+
+    def test_mro_entry_with_builtins(self):
+        tested = []
+        class A: ...
+        class C:
+            def __mro_entries__(self, bases):
+                tested.append(bases)
+                return (dict,)
+        c = C()
+        self.assertEqual(tested, [])
+        class D(A, c): ...
+        self.assertEqual(tested[-1], (A, c))
+        self.assertEqual(D.__bases__, (A, dict))
+        self.assertEqual(D.__orig_bases__, (A, c))
+        self.assertEqual(D.__mro__, (D, A, dict, object))
+
+    def test_mro_entry_with_builtins_2(self):
+        tested = []
+        class C:
+            def __mro_entries__(self, bases):
+                tested.append(bases)
+                return (C,)
+        c = C()
+        self.assertEqual(tested, [])
+        class D(c, dict): ...
+        self.assertEqual(tested[-1], (c, dict))
+        self.assertEqual(D.__bases__, (C, dict))
+        self.assertEqual(D.__orig_bases__, (c, dict))
+        self.assertEqual(D.__mro__, (D, C, dict, object))
+
+    def test_mro_entry_errors(self):
+        class C_too_many:
+            def __mro_entries__(self, bases, something, other):
+                return ()
+        c = C_too_many()
+        with self.assertRaises(TypeError):
+            class D(c): ...
+        class C_too_few:
+            def __mro_entries__(self):
+                return ()
+        d = C_too_few()
+        with self.assertRaises(TypeError):
+            class D(d): ...
+
+    def test_mro_entry_errors_2(self):
+        class C_not_callable:
+            __mro_entries__ = "Surprise!"
+        c = C_not_callable()
+        with self.assertRaises(TypeError):
+            class D(c): ...
+        class C_not_tuple:
+            def __mro_entries__(self):
+                return object
+        c = C_not_tuple()
+        with self.assertRaises(TypeError):
+            class D(c): ...
+
+    def test_mro_entry_metaclass(self):
+        meta_args = []
+        class Meta(type):
+            def __new__(mcls, name, bases, ns):
+                meta_args.extend([mcls, name, bases, ns])
+                return super().__new__(mcls, name, bases, ns)
+        class A: ...
+        class C:
+            def __mro_entries__(self, bases):
+                return (A,)
+        c = C()
+        class D(c, metaclass=Meta):
+            x = 1
+        self.assertEqual(meta_args[0], Meta)
+        self.assertEqual(meta_args[1], 'D')
+        self.assertEqual(meta_args[2], (A,))
+        self.assertEqual(meta_args[3]['x'], 1)
+        self.assertEqual(D.__bases__, (A,))
+        self.assertEqual(D.__orig_bases__, (c,))
+        self.assertEqual(D.__mro__, (D, A, object))
+        self.assertEqual(D.__class__, Meta)
+
+    def test_mro_entry_type_call(self):
+        # Substitution should _not_ happen in direct type call
+        class C:
+            def __mro_entries__(self, bases):
+                return ()
+        c = C()
+        with self.assertRaisesRegex(TypeError,
+                                    "MRO entry resolution; "
+                                    "use types.new_class()"):
+            type('Bad', (c,), {})
+
+
+class TestClassGetitem(unittest.TestCase):
+    def test_class_getitem(self):
+        getitem_args = []
+        class C:
+            def __class_getitem__(*args, **kwargs):
+                getitem_args.extend([args, kwargs])
+                return None
+        C[int, str]
+        self.assertEqual(getitem_args[0], (C, (int, str)))
+        self.assertEqual(getitem_args[1], {})
+
+    def test_class_getitem(self):
+        class C:
+            def __class_getitem__(cls, item):
+                return f'C[{item.__name__}]'
+        self.assertEqual(C[int], 'C[int]')
+        self.assertEqual(C[C], 'C[C]')
+
+    def test_class_getitem_inheritance(self):
+        class C:
+            def __class_getitem__(cls, item):
+                return f'{cls.__name__}[{item.__name__}]'
+        class D(C): ...
+        self.assertEqual(D[int], 'D[int]')
+        self.assertEqual(D[D], 'D[D]')
+
+    def test_class_getitem_inheritance_2(self):
+        class C:
+            def __class_getitem__(cls, item):
+                return 'Should not see this'
+        class D(C):
+            def __class_getitem__(cls, item):
+                return f'{cls.__name__}[{item.__name__}]'
+        self.assertEqual(D[int], 'D[int]')
+        self.assertEqual(D[D], 'D[D]')
+
+    def test_class_getitem_patched(self):
+        class C:
+            def __init_subclass__(cls):
+                def __class_getitem__(cls, item):
+                    return f'{cls.__name__}[{item.__name__}]'
+                cls.__class_getitem__ = __class_getitem__
+        class D(C): ...
+        self.assertEqual(D[int], 'D[int]')
+        self.assertEqual(D[D], 'D[D]')
+
+    def test_class_getitem_with_builtins(self):
+        class A(dict):
+            called_with = None
+
+            def __class_getitem__(cls, item):
+                cls.called_with = item
+        class B(A):
+            pass
+        self.assertIs(B.called_with, None)
+        B[int]
+        self.assertIs(B.called_with, int)
+
+    def test_class_getitem_errors(self):
+        class C_too_few:
+            def __class_getitem__(cls):
+                return None
+        with self.assertRaises(TypeError):
+            C_too_few[int]
+        class C_too_many:
+            def __class_getitem__(cls, one, two):
+                return None
+        with self.assertRaises(TypeError):
+            C_too_many[int]
+
+    def test_class_getitem_errors_2(self):
+        class C:
+            def __class_getitem__(cls, item):
+                return None
+        with self.assertRaises(TypeError):
+            C()[int]
+        class E: ...
+        e = E()
+        e.__class_getitem__ = lambda cls, item: 'This will not work'
+        with self.assertRaises(TypeError):
+            e[int]
+        class C_not_callable:
+            __class_getitem__ = "Surprise!"
+        with self.assertRaises(TypeError):
+            C_not_callable[int]
+
+    def test_class_getitem_metaclass(self):
+        class Meta(type):
+            def __class_getitem__(cls, item):
+                return f'{cls.__name__}[{item.__name__}]'
+        self.assertEqual(Meta[int], 'Meta[int]')
+
+    def test_class_getitem_metaclass_2(self):
+        class Meta(type):
+            def __getitem__(cls, item):
+                return 'from metaclass'
+        class C(metaclass=Meta):
+            def __class_getitem__(cls, item):
+                return 'from __class_getitem__'
+        self.assertEqual(C[int], 'from metaclass')
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/Lib/test/test_types.py b/Lib/test/test_types.py
index 28133a3..47488a6 100644
--- a/Lib/test/test_types.py
+++ b/Lib/test/test_types.py
@@ -844,6 +844,68 @@ class ClassCreationTests(unittest.TestCase):
         self.assertEqual(C.y, 1)
         self.assertEqual(C.z, 2)
 
+    def test_new_class_with_mro_entry(self):
+        class A: pass
+        class C:
+            def __mro_entries__(self, bases):
+                return (A,)
+        c = C()
+        D = types.new_class('D', (c,), {})
+        self.assertEqual(D.__bases__, (A,))
+        self.assertEqual(D.__orig_bases__, (c,))
+        self.assertEqual(D.__mro__, (D, A, object))
+
+    def test_new_class_with_mro_entry_none(self):
+        class A: pass
+        class B: pass
+        class C:
+            def __mro_entries__(self, bases):
+                return ()
+        c = C()
+        D = types.new_class('D', (A, c, B), {})
+        self.assertEqual(D.__bases__, (A, B))
+        self.assertEqual(D.__orig_bases__, (A, c, B))
+        self.assertEqual(D.__mro__, (D, A, B, object))
+
+    def test_new_class_with_mro_entry_error(self):
+        class A: pass
+        class C:
+            def __mro_entries__(self, bases):
+                return A
+        c = C()
+        with self.assertRaises(TypeError):
+            types.new_class('D', (c,), {})
+
+    def test_new_class_with_mro_entry_multiple(self):
+        class A1: pass
+        class A2: pass
+        class B1: pass
+        class B2: pass
+        class A:
+            def __mro_entries__(self, bases):
+                return (A1, A2)
+        class B:
+            def __mro_entries__(self, bases):
+                return (B1, B2)
+        D = types.new_class('D', (A(), B()), {})
+        self.assertEqual(D.__bases__, (A1, A2, B1, B2))
+
+    def test_new_class_with_mro_entry_multiple_2(self):
+        class A1: pass
+        class A2: pass
+        class A3: pass
+        class B1: pass
+        class B2: pass
+        class A:
+            def __mro_entries__(self, bases):
+                return (A1, A2, A3)
+        class B:
+            def __mro_entries__(self, bases):
+                return (B1, B2)
+        class C: pass
+        D = types.new_class('D', (A(), C, B()), {})
+        self.assertEqual(D.__bases__, (A1, A2, A3, C, B1, B2))
+
     # Many of the following tests are derived from test_descr.py
     def test_prepare_class(self):
         # Basic test of metaclass derivation
@@ -886,6 +948,28 @@ class ClassCreationTests(unittest.TestCase):
             class Bar(metaclass=BadMeta()):
                 pass
 
+    def test_resolve_bases(self):
+        class A: pass
+        class B: pass
+        class C:
+            def __mro_entries__(self, bases):
+                if A in bases:
+                    return ()
+                return (A,)
+        c = C()
+        self.assertEqual(types.resolve_bases(()), ())
+        self.assertEqual(types.resolve_bases((c,)), (A,))
+        self.assertEqual(types.resolve_bases((C,)), (C,))
+        self.assertEqual(types.resolve_bases((A, C)), (A, C))
+        self.assertEqual(types.resolve_bases((c, A)), (A,))
+        self.assertEqual(types.resolve_bases((A, c)), (A,))
+        x = (A,)
+        y = (C,)
+        z = (A, C)
+        t = (A, C, B)
+        for bases in [x, y, z, t]:
+            self.assertIs(types.resolve_bases(bases), bases)
+
     def test_metaclass_derivation(self):
         # issue1294232: correct metaclass calculation
         new_calls = []  # to check the order of __new__ calls
diff --git a/Lib/types.py b/Lib/types.py
index 336918f..c5976f3 100644
--- a/Lib/types.py
+++ b/Lib/types.py
@@ -60,10 +60,34 @@ del sys, _f, _g, _C, _c,                           # Not for export
 # Provide a PEP 3115 compliant mechanism for class creation
 def new_class(name, bases=(), kwds=None, exec_body=None):
     """Create a class object dynamically using the appropriate metaclass."""
-    meta, ns, kwds = prepare_class(name, bases, kwds)
+    resolved_bases = resolve_bases(bases)
+    meta, ns, kwds = prepare_class(name, resolved_bases, kwds)
     if exec_body is not None:
         exec_body(ns)
-    return meta(name, bases, ns, **kwds)
+    if resolved_bases is not bases:
+        ns['__orig_bases__'] = bases
+    return meta(name, resolved_bases, ns, **kwds)
+
+def resolve_bases(bases):
+    """Resolve MRO entries dynamically as specified by PEP 560."""
+    new_bases = list(bases)
+    updated = False
+    shift = 0
+    for i, base in enumerate(bases):
+        if isinstance(base, type):
+            continue
+        if not hasattr(base, "__mro_entries__"):
+            continue
+        new_base = base.__mro_entries__(bases)
+        updated = True
+        if not isinstance(new_base, tuple):
+            raise TypeError("__mro_entries__ must return a tuple")
+        else:
+            new_bases[i+shift:i+shift+1] = new_base
+            shift += len(new_base) - 1
+    if not updated:
+        return bases
+    return tuple(new_bases)
 
 def prepare_class(name, bases=(), kwds=None):
     """Call the __prepare__ method of the appropriate metaclass.
diff --git a/Misc/NEWS.d/next/Core and Builtins/2017-12-05-21-42-58.bpo-32226.G8fqb6.rst b/Misc/NEWS.d/next/Core and Builtins/2017-12-05-21-42-58.bpo-32226.G8fqb6.rst
new file mode 100644
index 0000000..97954fd
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2017-12-05-21-42-58.bpo-32226.G8fqb6.rst
@@ -0,0 +1,2 @@
+PEP 560: Add support for __mro_entries__ and __class_getitem__. Implemented
+by Ivan Levkivskyi.
diff --git a/Objects/abstract.c b/Objects/abstract.c
index 3cb7a32..0105c5d 100644
--- a/Objects/abstract.c
+++ b/Objects/abstract.c
@@ -168,6 +168,21 @@ PyObject_GetItem(PyObject *o, PyObject *key)
                               "be integer, not '%.200s'", key);
     }
 
+    if (PyType_Check(o)) {
+        PyObject *meth, *result, *stack[2] = {o, key};
+        _Py_IDENTIFIER(__class_getitem__);
+        meth = _PyObject_GetAttrId(o, &PyId___class_getitem__);
+        if (meth) {
+            result = _PyObject_FastCall(meth, stack, 2);
+            Py_DECREF(meth);
+            return result;
+        }
+        else if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
+            return NULL;
+        }
+        PyErr_Clear();
+    }
+
     return type_error("'%.200s' object is not subscriptable", o);
 }
 
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index 73f94e7..5403ecb 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -2377,6 +2377,27 @@ type_new(PyTypeObject *metatype, PyObject *args, PyObject *kwds)
         nbases = 1;
     }
     else {
+        _Py_IDENTIFIER(__mro_entries__);
+        for (i = 0; i < nbases; i++) {
+            tmp = PyTuple_GET_ITEM(bases, i);
+            if (PyType_Check(tmp)) {
+                continue;
+            }
+            tmp = _PyObject_GetAttrId(tmp, &PyId___mro_entries__);
+            if (tmp != NULL) {
+                PyErr_SetString(PyExc_TypeError,
+                                "type() doesn't support MRO entry resolution; "
+                                "use types.new_class()");
+                Py_DECREF(tmp);
+                return NULL;
+            }
+            else if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
+                PyErr_Clear();
+            }
+            else {
+                return NULL;
+            }
+        }
         /* Search the bases for the proper metatype to deal with this: */
         winner = _PyType_CalculateMetaclass(metatype, bases);
         if (winner == NULL) {
diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 23d7aa4..a363211 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -37,6 +37,7 @@ _Py_IDENTIFIER(__builtins__);
 _Py_IDENTIFIER(__dict__);
 _Py_IDENTIFIER(__prepare__);
 _Py_IDENTIFIER(__round__);
+_Py_IDENTIFIER(__mro_entries__);
 _Py_IDENTIFIER(encoding);
 _Py_IDENTIFIER(errors);
 _Py_IDENTIFIER(fileno);
@@ -49,12 +50,86 @@ _Py_IDENTIFIER(stderr);
 
 #include "clinic/bltinmodule.c.h"
 
+static PyObject*
+update_bases(PyObject *bases, PyObject *const *args, int nargs)
+{
+    int i, j;
+    PyObject *base, *meth, *new_base, *result, *new_bases = NULL;
+    PyObject *stack[1] = {bases};
+    assert(PyTuple_Check(bases));
+
+    for (i = 0; i < nargs; i++) {
+        base  = args[i];
+        if (PyType_Check(base)) {
+            if (new_bases) {
+                /* If we already have made a replacement, then we append every normal base,
+                   otherwise just skip it. */
+                if (PyList_Append(new_bases, base) < 0) {
+                    goto error;
+                }
+            }
+            continue;
+        }
+        meth = _PyObject_GetAttrId(base, &PyId___mro_entries__);
+        if (!meth) {
+            if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
+                goto error;
+            }
+            PyErr_Clear();
+            if (new_bases) {
+                if (PyList_Append(new_bases, base) < 0) {
+                    goto error;
+                }
+            }
+            continue;
+        }
+        new_base = _PyObject_FastCall(meth, stack, 1);
+        Py_DECREF(meth);
+        if (!new_base) {
+            goto error;
+        }
+        if (!PyTuple_Check(new_base)) {
+            PyErr_SetString(PyExc_TypeError,
+                            "__mro_entries__ must return a tuple");
+            Py_DECREF(new_base);
+            goto error;
+        }
+        if (!new_bases) {
+            /* If this is a first successful replacement, create new_bases list and
+               copy previously encountered bases. */
+            if (!(new_bases = PyList_New(i))) {
+                goto error;
+            }
+            for (j = 0; j < i; j++) {
+                base = args[j];
+                PyList_SET_ITEM(new_bases, j, base);
+                Py_INCREF(base);
+            }
+        }
+        j = PyList_GET_SIZE(new_bases);
+        if (PyList_SetSlice(new_bases, j, j, new_base) < 0) {
+            goto error;
+        }
+        Py_DECREF(new_base);
+    }
+    if (!new_bases) {
+        return bases;
+    }
+    result = PyList_AsTuple(new_bases);
+    Py_DECREF(new_bases);
+    return result;
+
+error:
+    Py_XDECREF(new_bases);
+    return NULL;
+}
+
 /* AC: cannot convert yet, waiting for *args support */
 static PyObject *
 builtin___build_class__(PyObject *self, PyObject **args, Py_ssize_t nargs,
                         PyObject *kwnames)
 {
-    PyObject *func, *name, *bases, *mkw, *meta, *winner, *prep, *ns;
+    PyObject *func, *name, *bases, *mkw, *meta, *winner, *prep, *ns, *orig_bases;
     PyObject *cls = NULL, *cell = NULL;
     int isclass = 0;   /* initialize to prevent gcc warning */
 
@@ -75,10 +150,16 @@ builtin___build_class__(PyObject *self, PyObject **args, Py_ssize_t nargs,
                         "__build_class__: name is not a string");
         return NULL;
     }
-    bases = _PyStack_AsTupleSlice(args, nargs, 2, nargs);
-    if (bases == NULL)
+    orig_bases = _PyStack_AsTupleSlice(args, nargs, 2, nargs);
+    if (orig_bases == NULL)
         return NULL;
 
+    bases = update_bases(orig_bases, args + 2, nargs - 2);
+    if (bases == NULL) {
+        Py_DECREF(orig_bases);
+        return NULL;
+    }
+
     if (kwnames == NULL) {
         meta = NULL;
         mkw = NULL;
@@ -171,6 +252,11 @@ builtin___build_class__(PyObject *self, PyObject **args, Py_ssize_t nargs,
                              NULL, 0, NULL, 0, NULL, 0, NULL,
                              PyFunction_GET_CLOSURE(func));
     if (cell != NULL) {
+        if (bases != orig_bases) {
+            if (PyMapping_SetItemString(ns, "__orig_bases__", orig_bases) < 0) {
+                goto error;
+            }
+        }
         PyObject *margs[3] = {name, bases, ns};
         cls = _PyObject_FastCallDict(meta, margs, 3, mkw);
         if (cls != NULL && PyType_Check(cls) && PyCell_Check(cell)) {
@@ -209,6 +295,9 @@ error:
     Py_DECREF(meta);
     Py_XDECREF(mkw);
     Py_DECREF(bases);
+    if (bases != orig_bases) {
+        Py_DECREF(orig_bases);
+    }
     return cls;
 }