implement __class_getitem__ from PEP 560

timfel · timfel · commit fefd8299acfa · 2019-04-01T12:36:29.000+02:00
diff --git a/graalpython/com.oracle.graal.python.test/src/tests/test_genericclass.py b/graalpython/com.oracle.graal.python.test/src/tests/test_genericclass.py
@@ -0,0 +1,116 @@
+# Copyright (c) 2018, 2019, Oracle and/or its affiliates.
+# Copyright (C) 1996-2017 Python Software Foundation
+#
+# Licensed under the PYTHON SOFTWARE FOUNDATION LICENSE VERSION 2
+import unittest
+
+
+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 'C[{0}]'.format(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 '{0}[{1}]'.format(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 '{0}[{1}]'.format(cls.__name__, item.__name__)
+        self.assertEqual(D[int], 'D[int]')
+        self.assertEqual(D[D], 'D[D]')
+
+    def test_class_getitem_classmethod(self):
+        class C:
+            @classmethod
+            def __class_getitem__(cls, item):
+                return '{0}[{1}]'.format(cls.__name__, item.__name__)
+        class D(C): ...
+        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 '{0}[{1}]'.format(cls.__name__, item.__name__)
+    #             cls.__class_getitem__ = classmethod(__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 '{0}[{1}]'.format(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')
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/BuiltinConstructors.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/BuiltinConstructors.java
@@ -129,6 +129,7 @@
 import com.oracle.graal.python.builtins.objects.type.TypeNodes.GetMroNode;
 import com.oracle.graal.python.builtins.objects.type.TypeNodes.GetNameNode;
 import com.oracle.graal.python.nodes.PGuards;
+import com.oracle.graal.python.nodes.SpecialMethodNames;
 import com.oracle.graal.python.nodes.attributes.GetAttributeNode.GetAnyAttributeNode;
 import com.oracle.graal.python.nodes.attributes.LookupAttributeInMRONode;
 import com.oracle.graal.python.nodes.attributes.LookupInheritedAttributeNode;
@@ -1896,6 +1897,23 @@ private Object typeMetaclass(String name, PTuple bases, PDict namespace, PythonA
                 Object value = entry.getValue();
                 if (__SLOTS__.equals(key)) {
                     slots = value;
+                } else if (SpecialMethodNames.__NEW__.equals(key)) {
+                    // TODO: see CPython: if it's a plain function, make it a
+                    // static function
+
+                    // tfel: this requires a little bit of refactoring on our
+                    // side that I don't want to do now
+                    pythonClass.setAttribute(key, value);
+                } else if (SpecialMethodNames.__INIT_SUBCLASS__.equals(key) ||
+                           SpecialMethodNames.__CLASS_GETITEM__.equals(key)) {
+                    // see CPython: Special-case __init_subclass__ and
+                    // __class_getitem__: if they are plain functions, make them
+                    // classmethods
+                    if (value instanceof PFunction) {
+                        pythonClass.setAttribute(key, factory().createClassmethod(value));
+                    } else {
+                        pythonClass.setAttribute(key, value);
+                    }
                 } else {
                     pythonClass.setAttribute(key, value);
                 }
@@ -1957,10 +1975,6 @@ private Object typeMetaclass(String name, PTuple bases, PDict namespace, PythonA
                 }
             }
 
-            // TODO: tfel special case __new__: if it's a plain function, make it a static function
-            // TODO: tfel Special-case __init_subclass__: if it's a plain function, make it a
-            // classmethod
-
             return pythonClass;
         }
 
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/SpecialMethodNames.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/SpecialMethodNames.java
@@ -161,6 +161,7 @@ public abstract class SpecialMethodNames {
     public static final String TOBYTES = "tobytes";
     public static final String DECODE = "decode";
     public static final String __SIZEOF__ = "__sizeof__";
+    public static final String __CLASS_GETITEM__ = "__class_getitem__";
 
     public static final String RICHCMP = "__truffle_richcompare__";
     public static final String TRUFFLE_SOURCE = "__truffle_source__";
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/subscript/GetItemNode.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/subscript/GetItemNode.java
@@ -25,22 +25,31 @@
  */
 package com.oracle.graal.python.nodes.subscript;
 
+import static com.oracle.graal.python.nodes.SpecialMethodNames.__CLASS_GETITEM__;
 import static com.oracle.graal.python.nodes.SpecialMethodNames.__GETITEM__;
+import static com.oracle.graal.python.runtime.exception.PythonErrorType.AttributeError;
 import static com.oracle.graal.python.runtime.exception.PythonErrorType.TypeError;
 
+import com.oracle.graal.python.builtins.objects.type.PythonAbstractClass;
 import com.oracle.graal.python.nodes.PRaiseNode;
+import com.oracle.graal.python.nodes.attributes.GetAttributeNode;
+import com.oracle.graal.python.nodes.call.CallNode;
 import com.oracle.graal.python.nodes.call.special.LookupAndCallBinaryNode;
 import com.oracle.graal.python.nodes.expression.BinaryOpNode;
 import com.oracle.graal.python.nodes.expression.ExpressionNode;
 import com.oracle.graal.python.nodes.frame.ReadNode;
+import com.oracle.graal.python.nodes.object.IsBuiltinClassProfile;
 import com.oracle.graal.python.nodes.statement.StatementNode;
+import com.oracle.graal.python.runtime.exception.PException;
 import com.oracle.truffle.api.CompilerDirectives;
+import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.NodeInfo;
 
 @NodeInfo(shortName = __GETITEM__)
 public abstract class GetItemNode extends BinaryOpNode implements ReadNode {
+    private static final String P_OBJECT_IS_NOT_SUBSCRIPTABLE = "'%p' object is not subscriptable";
 
     @Child private LookupAndCallBinaryNode callGetitemNode;
 
@@ -61,11 +70,39 @@ public Object doSpecialObject(Object primary, Object index) {
         if (callGetitemNode == null) {
             CompilerDirectives.transferToInterpreterAndInvalidate();
             callGetitemNode = insert(LookupAndCallBinaryNode.create(__GETITEM__, null, () -> new LookupAndCallBinaryNode.NotImplementedHandler() {
-                @Child private PRaiseNode raiseNode = PRaiseNode.create();
+                @CompilationFinal private IsBuiltinClassProfile isBuiltinClassProfile;
+                @Child private PRaiseNode raiseNode;
+                @Child private CallNode callClassGetItemNode;
+                @Child private GetAttributeNode getClassGetItemNode;
 
                 @Override
-                public Object execute(Object arg, @SuppressWarnings("unused") Object arg2) {
-                    throw raiseNode.raise(TypeError, "'%p' object is not subscriptable", arg);
+                public Object execute(Object arg, Object arg2) {
+                    if (arg instanceof PythonAbstractClass) {
+                        if (getClassGetItemNode == null) {
+                            CompilerDirectives.transferToInterpreterAndInvalidate();
+                            getClassGetItemNode = insert(GetAttributeNode.create(__CLASS_GETITEM__));
+                            isBuiltinClassProfile = IsBuiltinClassProfile.create();
+                        }
+                        Object classGetItem = null;
+                        try {
+                            classGetItem = getClassGetItemNode.executeObject(arg);
+                        } catch (PException e) {
+                            e.expect(AttributeError, isBuiltinClassProfile);
+                            // fall through to normal error handling
+                        }
+                        if (classGetItem != null) {
+                            if (callClassGetItemNode == null) {
+                                CompilerDirectives.transferToInterpreterAndInvalidate();
+                                callClassGetItemNode = insert(CallNode.create());
+                            }
+                            return callClassGetItemNode.execute(null, classGetItem, arg2);
+                        }
+                    }
+                    if (raiseNode == null) {
+                        CompilerDirectives.transferToInterpreterAndInvalidate();
+                        raiseNode = insert(PRaiseNode.create());
+                    }
+                    throw raiseNode.raise(TypeError, P_OBJECT_IS_NOT_SUBSCRIPTABLE, arg);
                 }
             }));
         }
diff --git a/mx.graalpython/copyrights/overrides b/mx.graalpython/copyrights/overrides
@@ -236,6 +236,7 @@ graalpython/com.oracle.graal.python.test/src/tests/test_getattribute-bimorphic-i
 graalpython/com.oracle.graal.python.test/src/tests/test_if-class-none.py,zippy.copyright
 graalpython/com.oracle.graal.python.test/src/tests/test_if.py,zippy.copyright
 graalpython/com.oracle.graal.python.test/src/tests/test_itertools.py,python.copyright
+graalpython/com.oracle.graal.python.test/src/tests/test_genericclass.py,python.copyright
 graalpython/com.oracle.graal.python.test/src/tests/test_list.py,python.copyright
 graalpython/com.oracle.graal.python.test/src/tests/test_builtin.py,python.copyright
 graalpython/com.oracle.graal.python.test/src/tests/test_mandelbrot3.py,benchmarks.copyright
