Generic work

Author: JelleZijlstra

Include/internal/pycore_intrinsics.h

@@ -10,6 +10,7 @@
 #define INTRINSIC_TYPEVAR 7
 #define INTRINSIC_PARAMSPEC 8
 #define INTRINSIC_TYPEVARTUPLE 9
+#define INTRINSIC_SUBSCRIPT_GENERIC 10
 
 #define MAX_INTRINSIC_1 9
 

Include/internal/pycore_typevarobject.h

@@ -13,10 +13,12 @@ extern PyTypeObject _PyTypeVarTuple_Type;
 extern PyTypeObject _PyParamSpec_Type;
 extern PyTypeObject _PyParamSpecArgs_Type;
 extern PyTypeObject _PyParamSpecKwargs_Type;
+extern PyTypeObject _PyGeneric_Type;
 
 extern PyObject *_Py_make_typevar(const char *, PyObject *);
 extern PyObject *_Py_make_paramspec(const char *);
 extern PyObject *_Py_make_typevartuple(const char *);
+extern PyObject *_Py_subscript_generic(PyObject *);
 
 #ifdef __cplusplus
 }

Lib/typing.py

@@ -1969,6 +1969,93 @@ def __init_subclass__(cls, *args, **kwargs):
         cls.__parameters__ = tuple(tvars)
 
 
+@_tp_cache
+def _generic_class_getitem(cls, params):
+    """Parameterizes a generic class.
+
+    At least, parameterizing a generic class is the *main* thing this method
+    does. For example, for some generic class `Foo`, this is called when we
+    do `Foo[int]` - there, with `cls=Foo` and `params=int`.
+
+    However, note that this method is also called when defining generic
+    classes in the first place with `class Foo(Generic[T]): ...`.
+    """
+    if not isinstance(params, tuple):
+        params = (params,)
+
+    params = tuple(_type_convert(p) for p in params)
+    if cls in (builtins.Generic, Protocol):
+        # Generic and Protocol can only be subscripted with unique type variables.
+        if not params:
+            raise TypeError(
+                f"Parameter list to {cls.__qualname__}[...] cannot be empty"
+            )
+        if not all(_is_typevar_like(p) for p in params):
+            raise TypeError(
+                f"Parameters to {cls.__name__}[...] must all be type variables "
+                f"or parameter specification variables.")
+        if len(set(params)) != len(params):
+            raise TypeError(
+                f"Parameters to {cls.__name__}[...] must all be unique")
+    else:
+        # Subscripting a regular Generic subclass.
+        for param in cls.__parameters__:
+            prepare = getattr(param, '__typing_prepare_subst__', None)
+            if prepare is not None:
+                params = prepare(cls, params)
+        _check_generic(cls, params, len(cls.__parameters__))
+
+        new_args = []
+        for param, new_arg in zip(cls.__parameters__, params):
+            if isinstance(param, TypeVarTuple):
+                new_args.extend(new_arg)
+            else:
+                new_args.append(new_arg)
+        params = tuple(new_args)
+
+    return _GenericAlias(cls, params,
+                            _paramspec_tvars=True)
+
+import builtins
+
+def _generic_init_subclass(cls, *args, **kwargs):
+    super(builtins.Generic, cls).__init_subclass__(*args, **kwargs)
+    tvars = []
+    if '__orig_bases__' in cls.__dict__:
+        error = builtins.Generic in cls.__orig_bases__
+    else:
+        error = (builtins.Generic in cls.__bases__ and
+                    cls.__name__ != 'Protocol' and
+                    type(cls) != _TypedDictMeta)
+    if error:
+        raise TypeError("Cannot inherit from plain Generic")
+    if '__orig_bases__' in cls.__dict__:
+        tvars = _collect_parameters(cls.__orig_bases__)
+        # Look for Generic[T1, ..., Tn].
+        # If found, tvars must be a subset of it.
+        # If not found, tvars is it.
+        # Also check for and reject plain Generic,
+        # and reject multiple Generic[...].
+        gvars = None
+        for base in cls.__orig_bases__:
+            if (isinstance(base, _GenericAlias) and
+                    base.__origin__ is builtins.Generic):
+                if gvars is not None:
+                    raise TypeError(
+                        "Cannot inherit from Generic[...] multiple types.")
+                gvars = base.__parameters__
+        if gvars is not None:
+            tvarset = set(tvars)
+            gvarset = set(gvars)
+            if not tvarset <= gvarset:
+                s_vars = ', '.join(str(t) for t in tvars if t not in gvarset)
+                s_args = ', '.join(str(g) for g in gvars)
+                raise TypeError(f"Some type variables ({s_vars}) are"
+                                f" not listed in Generic[{s_args}]")
+            tvars = gvars
+    cls.__parameters__ = tuple(tvars)
+
+
 class _TypingEllipsis:
     """Internal placeholder for ... (ellipsis)."""
 

Objects/typevarobject.c

@@ -11,6 +11,7 @@ class paramspec "paramspecobject *" "&_PyParamSpec_Type"
 class paramspecargs "paramspecargsobject *" "&_PyParamSpecArgs_Type"
 class paramspeckwargs "paramspeckwargsobject *" "&_PyParamSpecKwargs_Type"
 class typevartuple "typevartupleobject *" "&_PyTypeVarTuple_Type"
+class Generic "PyObject *" "&PyGeneric_Type
 [clinic start generated code]*/
 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=74cb9c15a049111b]*/
 
@@ -821,3 +822,111 @@ PyObject *_Py_make_paramspec(const char *name) {
 PyObject *_Py_make_typevartuple(const char *name) {
     return (PyObject *)typevartupleobject_alloc(name);
 }
+
+PyDoc_STRVAR(generic_doc,
+"Abstract base class for generic types.\n\
+\n\
+A generic type is typically declared by inheriting from\n\
+this class parameterized with one or more type variables.\n\
+For example, a generic mapping type might be defined as::\n\
+\n\
+    class Mapping(Generic[KT, VT]):\n\
+        def __getitem__(self, key: KT) -> VT:\n\
+            ...\n\
+        # Etc.\n\
+\n\
+This class can then be used as follows::\n\
+\n\
+    def lookup_name(mapping: Mapping[KT, VT], key: KT, default: VT) -> VT:\n\
+        try:\n\
+            return mapping[key]\n\
+        except KeyError:\n\
+            return default\n\
+");
+
+PyDoc_STRVAR(generic_class_getitem_doc,
+"Parameterizes a generic class.\n\
+\n\
+At least, parameterizing a generic class is the *main* thing this method\n\
+does. For example, for some generic class `Foo`, this is called when we\n\
+do `Foo[int]` - there, with `cls=Foo` and `params=int`.\n\
+\n\
+However, note that this method is also called when defining generic\n\
+classes in the first place with `class Foo(Generic[T]): ...`.\n\
+");
+
+static PyObject *
+call_typing_args_kwargs(const char *name, PyTypeObject *cls, PyObject *args, PyObject *kwargs)
+{
+    PyObject *typing = NULL, *func = NULL, *new_args = NULL;
+    typing = PyImport_ImportModule("typing");
+    if (typing == NULL) {
+        goto error;
+    }
+    func = PyObject_GetAttrString(typing, name);
+    if (func == NULL) {
+        goto error;
+    }
+    assert(PyTuple_Check(args));
+    Py_ssize_t nargs = PyTuple_GET_SIZE(args);
+    new_args = PyTuple_New(nargs + 1);
+    if (new_args == NULL) {
+        goto error;
+    }
+    PyTuple_SET_ITEM(new_args, 0, Py_NewRef((PyObject *)cls));
+    for (Py_ssize_t i = 0; i < nargs; i++) {
+        PyObject *arg = PyTuple_GET_ITEM(args, i);
+        PyTuple_SET_ITEM(new_args, i + 1, Py_NewRef(arg));
+    }
+    PyObject *result = PyObject_Call(func, new_args, kwargs);
+    Py_DECREF(func);
+    Py_DECREF(typing);
+    Py_DecRef(new_args);
+    return result;
+error:
+    Py_XDECREF(typing);
+    Py_XDECREF(func);
+    Py_XDECREF(new_args);
+    return NULL;
+}
+
+static PyObject *
+generic_init_subclass(PyTypeObject *cls, PyObject *args, PyObject *kwargs)
+{
+    return call_typing_args_kwargs("_generic_init_subclass", cls, args, kwargs);
+}
+
+static PyObject *
+generic_class_getitem(PyTypeObject *cls, PyObject *args, PyObject *kwargs)
+{
+    return call_typing_args_kwargs("_generic_class_getitem", cls, args, kwargs);
+}
+
+PyObject *
+_Py_subscript_generic(PyObject *params)
+{
+    PyObject *args = PyTuple_Pack(2, &_PyGeneric_Type, params);
+    if (args == NULL) {
+        return NULL;
+    }
+    return call_typing_func_object("_generic_class_getitem", args);
+}
+
+static PyMethodDef generic_methods[] = {
+    {"__class_getitem__", (PyCFunction)(void (*)(void))generic_class_getitem,
+     METH_VARARGS | METH_KEYWORDS | METH_CLASS,
+     generic_class_getitem_doc},
+    {"__init_subclass__", (PyCFunction)(void (*)(void))generic_init_subclass,
+     METH_VARARGS | METH_KEYWORDS | METH_CLASS,
+     PyDoc_STR("Function to initialize subclasses.")},
+    {NULL} /* Sentinel */
+};
+
+PyTypeObject _PyGeneric_Type = {
+    PyVarObject_HEAD_INIT(&PyType_Type, 0)
+    .tp_name = "typing.Generic",
+    .tp_basicsize = sizeof(PyObject),
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
+    .tp_doc = generic_doc,
+    .tp_methods = generic_methods,
+};

Python/bltinmodule.c

@@ -3096,6 +3096,7 @@ _PyBuiltin_Init(PyInterpreterState *interp)
     SETBUILTIN("TypeVar",               &_PyTypeVar_Type);
     SETBUILTIN("TypeVarTuple",          &_PyTypeVarTuple_Type);
     SETBUILTIN("ParamSpec",             &_PyParamSpec_Type);
+    SETBUILTIN("Generic",               &_PyGeneric_Type);
     debug = PyBool_FromLong(config->optimization_level == 0);
     if (PyDict_SetItemString(dict, "__debug__", debug) < 0) {
         Py_DECREF(debug);

Python/intrinsics.c

@@ -221,6 +221,12 @@ make_typevartuple(PyThreadState* unused, PyObject *v)
     return _Py_make_typevartuple(PyUnicode_AsUTF8(v));
 }
 
+static PyObject *
+subscript_generic(PyThreadState* unused, PyObject *v)
+{
+    return _Py_subscript_generic(v);
+}
+
 const instrinsic_func1
 _PyIntrinsics_UnaryFunctions[] = {
     [0] = no_intrinsic,
@@ -233,6 +239,7 @@ _PyIntrinsics_UnaryFunctions[] = {
     [INTRINSIC_TYPEVAR] = make_typevar,
     [INTRINSIC_PARAMSPEC] = make_paramspec,
     [INTRINSIC_TYPEVARTUPLE] = make_typevartuple,
+    [INTRINSIC_SUBSCRIPT_GENERIC] = subscript_generic,
 };
 
 

Python/symtable.c

@@ -1331,6 +1331,23 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
                                   DEF_LOCAL, LOCATION(s))) {
                 VISIT_QUIT(st, 0);
             }
+            for (Py_ssize_t i = 0; i < asdl_seq_LEN(s->v.ClassDef.typeparams); i++) {
+                typeparam_ty tp = asdl_seq_GET(s->v.ClassDef.typeparams, i);
+                PyObject *name;
+                switch (tp->kind) {
+                case TypeVar_kind:
+                    name = tp->v.TypeVar.name;
+                    break;
+                case ParamSpec_kind:
+                    name = tp->v.ParamSpec.name;
+                    break;
+                case TypeVarTuple_kind:
+                    name = tp->v.TypeVarTuple.name;
+                    break;
+                }
+                if (!symtable_add_def(st, name, USE, LOCATION(s)))
+                    VISIT_QUIT(st, 0);
+            }
         }
         VISIT_SEQ(st, stmt, s->v.ClassDef.body);
         st->st_private = tmp;