pythongh-104600: Make function.__type_params__ writable

JelleZijlstra · May 18, 2023 · 7c04199 · 7c04199
1 parent c5b670e
commit 7c04199
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 3 deletions.
diff --git a/Lib/functools.py b/Lib/functools.py
@@ -30,7 +30,7 @@
 # wrapper functions that can handle naive introspection
 
 WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__qualname__', '__doc__',
-                       '__annotations__')
+                       '__annotations__', '__type_params__')
 WRAPPER_UPDATES = ('__dict__',)
 def update_wrapper(wrapper,
                    wrapped,

diff --git a/Lib/test/test_funcattrs.py b/Lib/test/test_funcattrs.py
@@ -1,5 +1,6 @@
 import textwrap
 import types
+import typing
 import unittest
 
 
@@ -190,6 +191,20 @@ def test___qualname__(self):
         # __qualname__ must be a string
         self.cannot_set_attr(self.b, '__qualname__', 7, TypeError)
 
+    def test___type_params__(self):
+        def generic[T](): pass
+        def not_generic(): pass
+        T, = generic.__type_params__
+        self.assertIsInstance(T, typing.TypeVar)
+        self.assertEqual(generic.__type_params__, (T,))
+        self.assertEqual(not_generic.__type_params__, ())
+        with self.assertRaises(TypeError):
+            del not_generic.__type_params__
+        with self.assertRaises(TypeError):
+            not_generic.__type_params__ = 42
+        not_generic.__type_params__ = (T,)
+        self.assertEqual(not_generic.__type_params__, (T,))
+
     def test___code__(self):
         num_one, num_two = 7, 8
         def a(): pass

diff --git a/Lib/test/test_functools.py b/Lib/test/test_functools.py
@@ -617,7 +617,7 @@ def check_wrapper(self, wrapper, wrapped,
 
 
     def _default_update(self):
-        def f(a:'This is a new annotation'):
+        def f[T](a:'This is a new annotation'):
             """This is a test"""
             pass
         f.attr = 'This is also a test'
@@ -630,12 +630,14 @@ def wrapper(b:'This is the prior annotation'):
     def test_default_update(self):
         wrapper, f = self._default_update()
         self.check_wrapper(wrapper, f)
+        T, = f.__type_params__
         self.assertIs(wrapper.__wrapped__, f)
         self.assertEqual(wrapper.__name__, 'f')
         self.assertEqual(wrapper.__qualname__, f.__qualname__)
         self.assertEqual(wrapper.attr, 'This is also a test')
         self.assertEqual(wrapper.__annotations__['a'], 'This is a new annotation')
         self.assertNotIn('b', wrapper.__annotations__)
+        self.assertEqual(wrapper.__type_params__, (T,))
 
     @unittest.skipIf(sys.flags.optimize >= 2,
                      "Docstrings are omitted with -O2 and above")

diff --git a/Objects/funcobject.c b/Objects/funcobject.c
@@ -665,6 +665,20 @@ func_get_type_params(PyFunctionObject *op, void *Py_UNUSED(ignored))
     return Py_NewRef(op->func_typeparams);
 }
 
+static int
+func_set_type_params(PyFunctionObject *op, PyObject *value, void *Py_UNUSED(ignored))
+{
+    /* Not legal to del f.__type_params__ or to set it to anything
+     * other than a tuple object. */
+    if (value == NULL || !PyTuple_Check(value)) {
+        PyErr_SetString(PyExc_TypeError,
+                        "__type_params__ must be set to a tuple");
+        return -1;
+    }
+    Py_XSETREF(op->func_typeparams, Py_NewRef(value));
+    return 0;
+}
+
 PyObject *
 _Py_set_function_type_params(PyThreadState *Py_UNUSED(ignored), PyObject *func,
                              PyObject *type_params)
@@ -687,7 +701,8 @@ static PyGetSetDef func_getsetlist[] = {
     {"__dict__", PyObject_GenericGetDict, PyObject_GenericSetDict},
     {"__name__", (getter)func_get_name, (setter)func_set_name},
     {"__qualname__", (getter)func_get_qualname, (setter)func_set_qualname},
-    {"__type_params__", (getter)func_get_type_params, NULL},
+    {"__type_params__", (getter)func_get_type_params,
+     (setter)func_set_type_params},
     {NULL} /* Sentinel */
 };