[3.10] bpo-44490: Add __parameters__ and __getitem__ to types.Union (GH-26980) (GH-27207)

Co-authored-by: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com>
Co-authored-by: Guido van Rossum <gvanrossum@gmail.com>.
(cherry picked from commit c45fa1a5d9)

Co-authored-by: Yurii Karabas <1998uriyyo@gmail.com>
Co-authored-by: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com>

Include/internal/pycore_unionobject.h

@@ -12,6 +12,9 @@ PyAPI_FUNC(PyObject *) _Py_Union(PyObject *args);
 PyAPI_DATA(PyTypeObject) _Py_UnionType;
 PyAPI_FUNC(PyObject *) _Py_union_type_or(PyObject* self, PyObject* param);
 
+PyObject *_Py_subs_parameters(PyObject *, PyObject *, PyObject *, PyObject *);
+PyObject *_Py_make_parameters(PyObject *);
+
 #ifdef __cplusplus
 }
 #endif

Lib/test/test_types.py

@@ -711,6 +711,8 @@ class TypesTests(unittest.TestCase):
         TV = typing.TypeVar('T')
         assert TV | str == typing.Union[TV, str]
         assert str | TV == typing.Union[str, TV]
+        self.assertIs((int | TV)[int], int)
+        self.assertIs((TV | int)[int], int)
 
     def test_union_args(self):
         def check(arg, expected):
@@ -742,6 +744,17 @@ class TypesTests(unittest.TestCase):
                 check(x | None, (x, type(None)))
                 check(None | x, (type(None), x))
 
+    def test_union_parameter_chaining(self):
+        T = typing.TypeVar("T")
+        S = typing.TypeVar("S")
+
+        self.assertEqual((float | list[T])[int], float | list[int])
+        self.assertEqual(list[int | list[T]].__parameters__, (T,))
+        self.assertEqual(list[int | list[T]][str], list[int | list[str]])
+        self.assertEqual((list[T] | list[S]).__parameters__, (T, S))
+        self.assertEqual((list[T] | list[S])[int, T], list[int] | list[T])
+        self.assertEqual((list[T] | list[S])[int, int], list[int])
+
     def test_or_type_operator_with_forward(self):
         T = typing.TypeVar('T')
         ForwardAfter = T | 'Forward'

Misc/NEWS.d/next/Core and Builtins/2021-07-01-11-59-34.bpo-44490.xY80VR.rst

@@ -0,0 +1,2 @@
+Add ``__parameters__`` attribute and ``__getitem__``
+operator to ``types.Union``. Patch provided by Yurii Karabas.

Objects/genericaliasobject.c

@@ -198,8 +198,8 @@ tuple_add(PyObject *self, Py_ssize_t len, PyObject *item)
     return 0;
 }
 
-static PyObject *
-make_parameters(PyObject *args)
+PyObject *
+_Py_make_parameters(PyObject *args)
 {
     Py_ssize_t nargs = PyTuple_GET_SIZE(args);
     Py_ssize_t len = nargs;
@@ -294,18 +294,10 @@ subs_tvars(PyObject *obj, PyObject *params, PyObject **argitems)
     return obj;
 }
 
-static PyObject *
-ga_getitem(PyObject *self, PyObject *item)
+PyObject *
+_Py_subs_parameters(PyObject *self, PyObject *args, PyObject *parameters, PyObject *item)
 {
-    gaobject *alias = (gaobject *)self;
-    // do a lookup for __parameters__ so it gets populated (if not already)
-    if (alias->parameters == NULL) {
-        alias->parameters = make_parameters(alias->args);
-        if (alias->parameters == NULL) {
-            return NULL;
-        }
-    }
-    Py_ssize_t nparams = PyTuple_GET_SIZE(alias->parameters);
+    Py_ssize_t nparams = PyTuple_GET_SIZE(parameters);
     if (nparams == 0) {
         return PyErr_Format(PyExc_TypeError,
                             "There are no type variables left in %R",
@@ -320,32 +312,32 @@ ga_getitem(PyObject *self, PyObject *item)
                             nitems > nparams ? "many" : "few",
                             self);
     }
-    /* Replace all type variables (specified by alias->parameters)
+    /* Replace all type variables (specified by parameters)
        with corresponding values specified by argitems.
         t = list[T];          t[int]      -> newargs = [int]
         t = dict[str, T];     t[int]      -> newargs = [str, int]
         t = dict[T, list[S]]; t[str, int] -> newargs = [str, list[int]]
      */
-    Py_ssize_t nargs = PyTuple_GET_SIZE(alias->args);
+    Py_ssize_t nargs = PyTuple_GET_SIZE(args);
     PyObject *newargs = PyTuple_New(nargs);
     if (newargs == NULL) {
         return NULL;
     }
     for (Py_ssize_t iarg = 0; iarg < nargs; iarg++) {
-        PyObject *arg = PyTuple_GET_ITEM(alias->args, iarg);
+        PyObject *arg = PyTuple_GET_ITEM(args, iarg);
         int typevar = is_typevar(arg);
         if (typevar < 0) {
             Py_DECREF(newargs);
             return NULL;
         }
         if (typevar) {
-            Py_ssize_t iparam = tuple_index(alias->parameters, nparams, arg);
+            Py_ssize_t iparam = tuple_index(parameters, nparams, arg);
             assert(iparam >= 0);
             arg = argitems[iparam];
             Py_INCREF(arg);
         }
         else {
-            arg = subs_tvars(arg, alias->parameters, argitems);
+            arg = subs_tvars(arg, parameters, argitems);
             if (arg == NULL) {
                 Py_DECREF(newargs);
                 return NULL;
@@ -354,6 +346,26 @@ ga_getitem(PyObject *self, PyObject *item)
         PyTuple_SET_ITEM(newargs, iarg, arg);
     }
 
+    return newargs;
+}
+
+static PyObject *
+ga_getitem(PyObject *self, PyObject *item)
+{
+    gaobject *alias = (gaobject *)self;
+    // Populate __parameters__ if needed.
+    if (alias->parameters == NULL) {
+        alias->parameters = _Py_make_parameters(alias->args);
+        if (alias->parameters == NULL) {
+            return NULL;
+        }
+    }
+
+    PyObject *newargs = _Py_subs_parameters(self, alias->args, alias->parameters, item);
+    if (newargs == NULL) {
+        return NULL;
+    }
+
     PyObject *res = Py_GenericAlias(alias->origin, newargs);
 
     Py_DECREF(newargs);
@@ -550,7 +562,7 @@ ga_parameters(PyObject *self, void *unused)
 {
     gaobject *alias = (gaobject *)self;
     if (alias->parameters == NULL) {
-        alias->parameters = make_parameters(alias->args);
+        alias->parameters = _Py_make_parameters(alias->args);
         if (alias->parameters == NULL) {
             return NULL;
         }

Objects/unionobject.c

@@ -8,6 +8,7 @@
 typedef struct {
     PyObject_HEAD
     PyObject *args;
+    PyObject *parameters;
 } unionobject;
 
 static void
@@ -18,6 +19,7 @@ unionobject_dealloc(PyObject *self)
     _PyObject_GC_UNTRACK(self);
 
     Py_XDECREF(alias->args);
+    Py_XDECREF(alias->parameters);
     Py_TYPE(self)->tp_free(self);
 }
 
@@ -26,6 +28,7 @@ union_traverse(PyObject *self, visitproc visit, void *arg)
 {
     unionobject *alias = (unionobject *)self;
     Py_VISIT(alias->args);
+    Py_VISIT(alias->parameters);
     return 0;
 }
 
@@ -450,6 +453,53 @@ static PyMethodDef union_methods[] = {
         {"__subclasscheck__", union_subclasscheck, METH_O},
         {0}};
 
+
+static PyObject *
+union_getitem(PyObject *self, PyObject *item)
+{
+    unionobject *alias = (unionobject *)self;
+    // Populate __parameters__ if needed.
+    if (alias->parameters == NULL) {
+        alias->parameters = _Py_make_parameters(alias->args);
+        if (alias->parameters == NULL) {
+            return NULL;
+        }
+    }
+
+    PyObject *newargs = _Py_subs_parameters(self, alias->args, alias->parameters, item);
+    if (newargs == NULL) {
+        return NULL;
+    }
+
+    PyObject *res = _Py_Union(newargs);
+
+    Py_DECREF(newargs);
+    return res;
+}
+
+static PyMappingMethods union_as_mapping = {
+    .mp_subscript = union_getitem,
+};
+
+static PyObject *
+union_parameters(PyObject *self, void *Py_UNUSED(unused))
+{
+    unionobject *alias = (unionobject *)self;
+    if (alias->parameters == NULL) {
+        alias->parameters = _Py_make_parameters(alias->args);
+        if (alias->parameters == NULL) {
+            return NULL;
+        }
+    }
+    Py_INCREF(alias->parameters);
+    return alias->parameters;
+}
+
+static PyGetSetDef union_properties[] = {
+    {"__parameters__", union_parameters, (setter)NULL, "Type variables in the types.Union.", NULL},
+    {0}
+};
+
 static PyNumberMethods union_as_number = {
         .nb_or = _Py_union_type_or, // Add __or__ function
 };
@@ -471,8 +521,10 @@ PyTypeObject _Py_UnionType = {
     .tp_members = union_members,
     .tp_methods = union_methods,
     .tp_richcompare = union_richcompare,
+    .tp_as_mapping = &union_as_mapping,
     .tp_as_number = &union_as_number,
     .tp_repr = union_repr,
+    .tp_getset = union_properties,
 };
 
 PyObject *
@@ -516,6 +568,7 @@ _Py_Union(PyObject *args)
         return NULL;
     }
 
+    result->parameters = NULL;
     result->args = args;
     _PyObject_GC_TRACK(result);
     return (PyObject*)result;