bpo-44653: Support typing.Union in parameter substitution of the union type

serhiy-storchaka · serhiy-storchaka · commit a74e69f53808 · 2021-07-18T19:50:24.000+03:00
diff --git a/Lib/test/test_types.py b/Lib/test/test_types.py
@@ -772,6 +772,21 @@ def test_union_parameter_chaining(self):
         self.assertEqual((list[T] | list[S])[int, T], list[int] | list[T])
         self.assertEqual((list[T] | list[S])[int, int], list[int])
 
+    def test_union_parameter_substitution_union(self):
+        T = typing.TypeVar("T")
+        res = (int | T)[str | list]
+        self.assertEqual(res, int | str | list)
+        self.assertIsInstance(res, types.Union)
+        res = (int | T)[str | int]
+        self.assertEqual(res, int | str)
+        self.assertIsInstance(res, types.Union)
+        res = (int | T)[typing.Union[str, list]]
+        self.assertEqual(res, int | str | list)
+        self.assertIsInstance(res, types.Union)
+        res = (int | T)[typing.Union[str, int]]
+        self.assertEqual(res, int | str)
+        self.assertIsInstance(res, types.Union)
+
     def test_union_parameter_substitution_errors(self):
         T = typing.TypeVar("T")
         x = int | T
diff --git a/Misc/NEWS.d/next/Core and Builtins/2021-07-18-19-32-38.bpo-44653.lu99mE.rst b/Misc/NEWS.d/next/Core and Builtins/2021-07-18-19-32-38.bpo-44653.lu99mE.rst
@@ -0,0 +1 @@
+Support ``typing.Union`` in parameter substitution of the union type.
diff --git a/Objects/unionobject.c b/Objects/unionobject.c
@@ -273,6 +273,12 @@ is_new_type(PyObject *obj)
     return is_typing_module(obj);
 }
 
+static int
+is_typing_union(PyObject *obj)
+{
+    return is_typing_name(obj, "_UnionGenericAlias");
+}
+
 // Emulates short-circuiting behavior of the ``||`` operator
 // while also checking negative values.
 #define CHECK_RES(res) { \
@@ -429,6 +435,19 @@ static PyMethodDef union_methods[] = {
         {0}};
 
 
+static PyObject *
+from_typing_union(PyObject *obj)
+{
+    _Py_IDENTIFIER(__args__);
+    PyObject *args = _PyObject_GetAttrId(obj, &PyId___args__);
+    if (args == NULL) {
+        return NULL;
+    }
+    PyObject *result = make_union(args);
+    Py_DECREF(args);
+    return result;
+}
+
 static PyObject *
 union_getitem(PyObject *self, PyObject *item)
 {
@@ -447,16 +466,34 @@ union_getitem(PyObject *self, PyObject *item)
     }
 
     // Check arguments are unionable.
+    assert(Py_REFCNT(newargs) == 1);
     Py_ssize_t nargs = PyTuple_GET_SIZE(newargs);
     for (Py_ssize_t iarg = 0; iarg < nargs; iarg++) {
         PyObject *arg = PyTuple_GET_ITEM(newargs, iarg);
-        int is_arg_unionable = is_unionable(arg);
-        if (is_arg_unionable <= 0) {
-            Py_DECREF(newargs);
-            if (is_arg_unionable == 0) {
+        int r = is_unionable(arg);
+        if (r == 0) {
+            r = is_typing_union(arg);
+            if (r == 0) {
                 PyErr_Format(PyExc_TypeError,
                              "Each union argument must be a type, got %.100R", arg);
+                Py_DECREF(newargs);
+                return NULL;
+            }
+            if (r > 0) {
+                // Replace typing.Union with types.Union.
+                PyObject *newarg = from_typing_union(arg);
+                if (newarg == NULL) {
+                    Py_DECREF(newargs);
+                    return NULL;
+                }
+                assert(Py_REFCNT(newargs) == 1);
+                assert(PyTuple_GET_ITEM(newargs, iarg) == arg);
+                PyTuple_SET_ITEM(newargs, iarg, newarg);
+                Py_DECREF(arg);
             }
+        }
+        if (r < 0) {
+            Py_DECREF(newargs);
             return NULL;
         }
     }

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Support ``typing.Union`` in parameter substitution of the union type.