bpo-44676: Add ability to serialize types.Union (GH-27244)

uriyyo · web-flow · commit fe13f0b0f696 · 2021-07-23T11:47:00.000+02:00
diff --git a/Lib/test/test_types.py b/Lib/test/test_types.py
@@ -13,6 +13,9 @@
 import weakref
 import typing
 
+
+T = typing.TypeVar("T")
+
 class Example:
     pass
 
@@ -802,6 +805,38 @@ def eq(actual, expected):
         eq(x[NT], int | NT | bytes)
         eq(x[S], int | S | bytes)
 
+    def test_union_pickle(self):
+        alias = list[T] | int
+        s = pickle.dumps(alias)
+        loaded = pickle.loads(s)
+        self.assertEqual(alias, loaded)
+        self.assertEqual(alias.__args__, loaded.__args__)
+        self.assertEqual(alias.__parameters__, loaded.__parameters__)
+
+    def test_union_from_args(self):
+        with self.assertRaisesRegex(
+                TypeError,
+                r"^Each union argument must be a type, got 1$",
+        ):
+            types.Union._from_args((1,))
+
+        with self.assertRaisesRegex(
+                TypeError,
+                r"Union._from_args\(\) argument 'args' must be tuple, not int$",
+        ):
+            types.Union._from_args(1)
+
+        with self.assertRaisesRegex(ValueError, r"args must be not empty"):
+            types.Union._from_args(())
+
+        alias = types.Union._from_args((int, str, T))
+
+        self.assertEqual(alias.__args__, (int, str, T))
+        self.assertEqual(alias.__parameters__, (T,))
+
+        result = types.Union._from_args((int,))
+        self.assertIs(int, result)
+
     def test_union_parameter_substitution_errors(self):
         T = typing.TypeVar("T")
         x = int | T
diff --git a/Lib/typing.py b/Lib/typing.py
@@ -329,7 +329,7 @@ def _eval_type(t, globalns, localns, recursive_guard=frozenset()):
         if isinstance(t, GenericAlias):
             return GenericAlias(t.__origin__, ev_args)
         if isinstance(t, types.Union):
-            return functools.reduce(operator.or_, ev_args)
+            return types.Union._from_args(ev_args)
         else:
             return t.copy_with(ev_args)
     return t
@@ -1808,7 +1808,7 @@ def _strip_annotations(t):
         stripped_args = tuple(_strip_annotations(a) for a in t.__args__)
         if stripped_args == t.__args__:
             return t
-        return functools.reduce(operator.or_, stripped_args)
+        return types.Union._from_args(stripped_args)
 
     return t
 
diff --git a/Misc/NEWS.d/next/Core and Builtins/2021-07-19-19-53-46.bpo-44676.WgIMvh.rst b/Misc/NEWS.d/next/Core and Builtins/2021-07-19-19-53-46.bpo-44676.WgIMvh.rst
@@ -0,0 +1,2 @@
+Add ability to serialise ``types.Union`` objects. Patch provided by Yurii
+Karabas.
diff --git a/Objects/unionobject.c b/Objects/unionobject.c
@@ -299,6 +299,24 @@ is_unionable(PyObject *obj)
     return 0;
 }
 
+static int
+is_args_unionable(PyObject *args)
+{
+    Py_ssize_t nargs = PyTuple_GET_SIZE(args);
+    for (Py_ssize_t iarg = 0; iarg < nargs; iarg++) {
+        PyObject *arg = PyTuple_GET_ITEM(args, iarg);
+        int is_arg_unionable = is_unionable(arg);
+        if (is_arg_unionable <= 0) {
+            if (is_arg_unionable == 0) {
+                PyErr_Format(PyExc_TypeError,
+                             "Each union argument must be a type, got %.100R", arg);
+            }
+            return 0;
+        }
+    }
+    return 1;
+}
+
 PyObject *
 _Py_union_type_or(PyObject* self, PyObject* other)
 {
@@ -418,14 +436,47 @@ union_repr(PyObject *self)
     return NULL;
 }
 
+static PyObject *
+union_reduce(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    unionobject *alias = (unionobject *)self;
+    PyObject* from_args = PyObject_GetAttrString(self, "_from_args");
+    if (from_args == NULL) {
+        return NULL;
+    }
+
+    return Py_BuildValue("O(O)", from_args, alias->args);
+}
+
 static PyMemberDef union_members[] = {
         {"__args__", T_OBJECT, offsetof(unionobject, args), READONLY},
         {0}
 };
 
+static PyObject *
+union_from_args(PyObject *cls, PyObject *args)
+{
+    if (!PyTuple_CheckExact(args)) {
+        _PyArg_BadArgument("Union._from_args", "argument 'args'", "tuple", args);
+        return NULL;
+    }
+    if (!PyTuple_GET_SIZE(args)) {
+        PyErr_SetString(PyExc_ValueError, "args must be not empty");
+        return NULL;
+    }
+
+    if (is_args_unionable(args) <= 0) {
+        return NULL;
+    }
+
+    return make_union(args);
+}
+
 static PyMethodDef union_methods[] = {
+        {"_from_args", union_from_args, METH_O | METH_CLASS},
         {"__instancecheck__", union_instancecheck, METH_O},
         {"__subclasscheck__", union_subclasscheck, METH_O},
+        {"__reduce__", union_reduce, METH_NOARGS},
         {0}};
 
 

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	+Add ability to serialise ``types.Union`` objects. Patch provided by Yurii
	`2`	`+Karabas.`