gh-129250: allow pickle instances of generic classes

python · tom-pytel · Jan 29, 2025 · Jan 29, 2025 · Jan 29, 2025 · Jan 29, 2025
commit 8dbbbcd356b8e9e0d83620a91f056c12cb98d9dc
diff --git a/Include/internal/pycore_typevarobject.h b/Include/internal/pycore_typevarobject.h
@@ -17,6 +17,7 @@ extern PyObject *_Py_set_typeparam_default(PyThreadState *, PyObject *, PyObject
 extern int _Py_initialize_generic(PyInterpreterState *);
 extern void _Py_clear_generic_types(PyInterpreterState *);
 extern int _Py_typing_type_repr(PyUnicodeWriter *, PyObject *);
+extern int _Py_set_type_params_owner_maybe(PyObject *, PyObject *);

 extern PyTypeObject _PyTypeAlias_Type;
 extern PyTypeObject _PyNoDefault_Type;

@@ -1183,11 +1183,14 @@ def func1[X](x: X) -> X: ...
 def func2[X, Y](x: X | Y) -> X | Y: ...
 def func3[X, *Y, **Z](x: X, y: tuple[*Y], z: Z) -> X: ...
 def func4[X: int, Y: (bytes, str)](x: X, y: Y) -> X | Y: ...
+def func3b[X, *Y, **Z]() -> X: return Class3[X, Y, Z]()
+def func5[Baz](): return Class5[Baz]()

 class Class1[X]: ...
 class Class2[X, Y]: ...
 class Class3[X, *Y, **Z]: ...
 class Class4[X: int, Y: (bytes, str)]: ...
+class Class5(Generic[T]): pass


 class TypeParamsPickleTest(unittest.TestCase):
@@ -1240,6 +1243,32 @@ def test_pickling_classes(self):
                    # but class check is good enough:
                    self.assertIsInstance(pickle.loads(pickled), real_class)

+    def test_pickling_anonymous_typeparams(self):
+        # see gh-129250
+        thing = func5()
+        pickled = pickle.dumps(thing)
+        unpickled = pickle.loads(pickled)
+        self.assertIs(unpickled.__orig_class__, thing.__orig_class__)
+        self.assertIs(unpickled.__orig_class__.__args__[0],
+                      func5.__type_params__[0])
+
+        thing = func3b()
+        pickled = pickle.dumps(thing)
+        unpickled = pickle.loads(pickled)
+        self.assertIs(unpickled.__orig_class__, thing.__orig_class__)
+        self.assertIs(unpickled.__orig_class__.__args__[0],
+                      func3b.__type_params__[0])
+        self.assertIs(unpickled.__orig_class__.__args__[1],
+                      func3b.__type_params__[1])
+        self.assertIs(unpickled.__orig_class__.__args__[2],
+                      func3b.__type_params__[2])
+
+        for i in range(3):
+            thing = Class3.__type_params__[i]
+            pickled = pickle.dumps(thing)
+            unpickled = pickle.loads(pickled)
+            self.assertIs(unpickled, thing)
+

 class TypeParamsWeakRefTest(unittest.TestCase):
    def test_weakrefs(self):

diff --git a/Lib/typing.py b/Lib/typing.py
@@ -502,6 +502,10 @@ def _eval_type(t, globalns, localns, type_params=_sentinel, *, recursive_guard=f
    return t


+def _restore_anonymous_typeparam(owner, index):
+    return owner.__type_params__[index]
+
+
 class _Final:
    """Mixin to prohibit subclassing."""


diff --git a/Objects/funcobject.c b/Objects/funcobject.c
@@ -2,11 +2,12 @@
 /* Function object implementation */

 #include "Python.h"
-#include "pycore_dict.h"          // _Py_INCREF_DICT()
-#include "pycore_long.h"          // _PyLong_GetOne()
-#include "pycore_modsupport.h"    // _PyArg_NoKeywords()
-#include "pycore_object.h"        // _PyObject_GC_UNTRACK()
-#include "pycore_pyerrors.h"      // _PyErr_Occurred()
+#include "pycore_dict.h"           // _Py_INCREF_DICT()
+#include "pycore_long.h"           // _PyLong_GetOne()
+#include "pycore_modsupport.h"     // _PyArg_NoKeywords()
+#include "pycore_object.h"         // _PyObject_GC_UNTRACK()
+#include "pycore_pyerrors.h"       // _PyErr_Occurred()
+#include "pycore_typevarobject.h"  // _Py_set_type_params_owner_maybe()


 static const char *
@@ -926,6 +927,7 @@ _Py_set_function_type_params(PyThreadState *Py_UNUSED(ignored), PyObject *func,
    assert(PyFunction_Check(func));
    assert(PyTuple_Check(type_params));
    PyFunctionObject *f = (PyFunctionObject *)func;
+    _Py_set_type_params_owner_maybe(type_params, func);
    Py_XSETREF(f->func_typeparams, Py_NewRef(type_params));
    return Py_NewRef(func);
 }