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

Author: uriyyo

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

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
 

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.

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}};