gh-104549: Set __module__ on TypeAliasType (#104550)

Author: JelleZijlstra

Include/internal/pycore_global_objects.h

@@ -75,7 +75,6 @@ struct _Py_interp_cached_objects {
     PyTypeObject *paramspec_type;
     PyTypeObject *paramspecargs_type;
     PyTypeObject *paramspeckwargs_type;
-    PyTypeObject *typealias_type;
 };
 
 #define _Py_INTERP_STATIC_OBJECT(interp, NAME) \

Include/internal/pycore_typevarobject.h

@@ -16,6 +16,8 @@ extern PyObject *_Py_subscript_generic(PyThreadState *, PyObject *);
 extern int _Py_initialize_generic(PyInterpreterState *);
 extern void _Py_clear_generic_types(PyInterpreterState *);
 
+extern PyTypeObject _PyTypeAlias_Type;
+
 #ifdef __cplusplus
 }
 #endif

Lib/test/mod_generics_cache.py

@@ -1,6 +1,6 @@
 """Module for testing the behavior of generics across different modules."""
 
-from typing import TypeVar, Generic, Optional
+from typing import TypeVar, Generic, Optional, TypeAliasType
 
 default_a: Optional['A'] = None
 default_b: Optional['B'] = None
@@ -19,3 +19,6 @@ class A(Generic[T]):
     my_inner_a1: 'B.A'
     my_inner_a2: A
     my_outer_a: 'A'  # unless somebody calls get_type_hints with localns=B.__dict__
+
+type Alias = int
+OldStyle = TypeAliasType("OldStyle", int)

Lib/test/test_type_aliases.py

@@ -1,6 +1,8 @@
+import pickle
 import types
 import unittest
 from test.support import check_syntax_error, run_code
+from test import mod_generics_cache
 
 from typing import Callable, TypeAliasType, TypeVar, get_args
 
@@ -155,19 +157,22 @@ def test_basic(self):
         self.assertEqual(TA.__name__, "TA")
         self.assertIs(TA.__value__, int)
         self.assertEqual(TA.__type_params__, ())
+        self.assertEqual(TA.__module__, __name__)
 
     def test_generic(self):
         T = TypeVar("T")
         TA = TypeAliasType("TA", list[T], type_params=(T,))
         self.assertEqual(TA.__name__, "TA")
         self.assertEqual(TA.__value__, list[T])
         self.assertEqual(TA.__type_params__, (T,))
+        self.assertEqual(TA.__module__, __name__)
 
     def test_keywords(self):
         TA = TypeAliasType(name="TA", value=int)
         self.assertEqual(TA.__name__, "TA")
         self.assertIs(TA.__value__, int)
         self.assertEqual(TA.__type_params__, ())
+        self.assertEqual(TA.__module__, __name__)
 
     def test_errors(self):
         with self.assertRaises(TypeError):
@@ -202,3 +207,18 @@ def test_union(self):
         union3 = list[range] | Alias1
         self.assertIsInstance(union3, types.UnionType)
         self.assertEqual(get_args(union3), (list[range], Alias1))
+
+    def test_module(self):
+        self.assertEqual(TypeAliasType.__module__, "typing")
+        type Alias = int
+        self.assertEqual(Alias.__module__, __name__)
+        self.assertEqual(mod_generics_cache.Alias.__module__,
+                         mod_generics_cache.__name__)
+        self.assertEqual(mod_generics_cache.OldStyle.__module__,
+                         mod_generics_cache.__name__)
+
+    def test_pickling(self):
+        pickled = pickle.dumps(mod_generics_cache.Alias)
+        self.assertIs(pickle.loads(pickled), mod_generics_cache.Alias)
+        pickled = pickle.dumps(mod_generics_cache.OldStyle)
+        self.assertIs(pickle.loads(pickled), mod_generics_cache.OldStyle)

Modules/_typingmodule.c

@@ -6,6 +6,7 @@
 
 #include "Python.h"
 #include "internal/pycore_interp.h"
+#include "internal/pycore_typevarobject.h"
 #include "clinic/_typingmodule.c.h"
 
 /*[clinic input]
@@ -56,9 +57,11 @@ _typing_exec(PyObject *m)
     EXPORT_TYPE("ParamSpec", paramspec_type);
     EXPORT_TYPE("ParamSpecArgs", paramspecargs_type);
     EXPORT_TYPE("ParamSpecKwargs", paramspeckwargs_type);
-    EXPORT_TYPE("TypeAliasType", typealias_type);
     EXPORT_TYPE("Generic", generic_type);
 #undef EXPORT_TYPE
+    if (PyModule_AddObjectRef(m, "TypeAliasType", (PyObject *)&_PyTypeAlias_Type) < 0) {
+        return -1;
+    }
     return 0;
 }
 

Objects/object.c

@@ -14,7 +14,7 @@
 #include "pycore_pymem.h"         // _PyMem_IsPtrFreed()
 #include "pycore_pystate.h"       // _PyThreadState_GET()
 #include "pycore_symtable.h"      // PySTEntry_Type
-#include "pycore_typevarobject.h" // _PyTypeVar_Type etc., _Py_initialize_generic
+#include "pycore_typevarobject.h" // _PyTypeAlias_Type, _Py_initialize_generic
 #include "pycore_typeobject.h"    // _PyBufferWrapper_Type
 #include "pycore_unionobject.h"   // _PyUnion_Type
 #include "pycore_interpreteridobject.h"  // _PyInterpreterID_Type
@@ -2112,6 +2112,7 @@ static PyTypeObject* static_types[] = {
     &_PyWeakref_CallableProxyType,
     &_PyWeakref_ProxyType,
     &_PyWeakref_RefType,
+    &_PyTypeAlias_Type,
 
     // subclasses: _PyTypes_FiniTypes() deallocates them before their base
     // class

Objects/typevarobject.c

@@ -48,6 +48,7 @@ typedef struct {
     PyObject *type_params;
     PyObject *compute_value;
     PyObject *value;
+    PyObject *module;
 } typealiasobject;
 
 #include "clinic/typevarobject.c.h"
@@ -1252,6 +1253,7 @@ typealias_dealloc(PyObject *self)
     Py_XDECREF(ta->type_params);
     Py_XDECREF(ta->compute_value);
     Py_XDECREF(ta->value);
+    Py_XDECREF(ta->module);
     Py_TYPE(self)->tp_free(self);
     Py_DECREF(tp);
 }
@@ -1309,26 +1311,41 @@ typealias_type_params(PyObject *self, void *unused)
     return Py_NewRef(ta->type_params);
 }
 
+static PyObject *
+typealias_module(PyObject *self, void *unused)
+{
+    typealiasobject *ta = (typealiasobject *)self;
+    if (ta->module != NULL) {
+        return Py_NewRef(ta->module);
+    }
+    if (ta->compute_value != NULL) {
+        // PyFunction_GetModule() returns a borrowed reference
+        return Py_NewRef(PyFunction_GetModule(ta->compute_value));
+    }
+    Py_RETURN_NONE;
+}
+
 static PyGetSetDef typealias_getset[] = {
     {"__parameters__", typealias_parameters, (setter)NULL, NULL, NULL},
     {"__type_params__", typealias_type_params, (setter)NULL, NULL, NULL},
     {"__value__", typealias_value, (setter)NULL, NULL, NULL},
+    {"__module__", typealias_module, (setter)NULL, NULL, NULL},
     {0}
 };
 
 static typealiasobject *
 typealias_alloc(PyObject *name, PyObject *type_params, PyObject *compute_value,
-                PyObject *value)
+                PyObject *value, PyObject *module)
 {
-    PyTypeObject *tp = PyInterpreterState_Get()->cached_objects.typealias_type;
-    typealiasobject *ta = PyObject_GC_New(typealiasobject, tp);
+    typealiasobject *ta = PyObject_GC_New(typealiasobject, &_PyTypeAlias_Type);
     if (ta == NULL) {
         return NULL;
     }
     ta->name = Py_NewRef(name);
     ta->type_params = Py_IsNone(type_params) ? NULL : Py_XNewRef(type_params);
     ta->compute_value = Py_XNewRef(compute_value);
     ta->value = Py_XNewRef(value);
+    ta->module = Py_XNewRef(module);
     _PyObject_GC_TRACK(ta);
     return ta;
 }
@@ -1339,6 +1356,7 @@ typealias_traverse(typealiasobject *self, visitproc visit, void *arg)
     Py_VISIT(self->type_params);
     Py_VISIT(self->compute_value);
     Py_VISIT(self->value);
+    Py_VISIT(self->module);
     return 0;
 }
 
@@ -1348,6 +1366,7 @@ typealias_clear(typealiasobject *self)
     Py_CLEAR(self->type_params);
     Py_CLEAR(self->compute_value);
     Py_CLEAR(self->value);
+    Py_CLEAR(self->module);
     return 0;
 }
 
@@ -1401,7 +1420,14 @@ typealias_new_impl(PyTypeObject *type, PyObject *name, PyObject *value,
         PyErr_SetString(PyExc_TypeError, "type_params must be a tuple");
         return NULL;
     }
-    return (PyObject *)typealias_alloc(name, type_params, NULL, value);
+    PyObject *module = caller();
+    if (module == NULL) {
+        return NULL;
+    }
+    PyObject *ta = (PyObject *)typealias_alloc(name, type_params, NULL, value,
+                                               module);
+    Py_DECREF(module);
+    return ta;
 }
 
 PyDoc_STRVAR(typealias_doc,
@@ -1412,28 +1438,32 @@ Type aliases are created through the type statement:\n\
   type Alias = int\n\
 ");
 
-static PyType_Slot typealias_slots[] = {
-    {Py_tp_doc, (void *)typealias_doc},
-    {Py_tp_members, typealias_members},
-    {Py_tp_methods, typealias_methods},
-    {Py_tp_getset, typealias_getset},
-    {Py_mp_subscript, typealias_subscript},
-    {Py_tp_dealloc, typealias_dealloc},
-    {Py_tp_alloc, PyType_GenericAlloc},
-    {Py_tp_new, typealias_new},
-    {Py_tp_free, PyObject_GC_Del},
-    {Py_tp_traverse, (traverseproc)typealias_traverse},
-    {Py_tp_clear, (inquiry)typealias_clear},
-    {Py_tp_repr, typealias_repr},
-    {Py_nb_or, _Py_union_type_or},
-    {0, 0},
+static PyNumberMethods typealias_as_number = {
+    .nb_or = _Py_union_type_or,
+};
+
+static PyMappingMethods typealias_as_mapping = {
+    .mp_subscript = typealias_subscript,
 };
 
-PyType_Spec typealias_spec = {
-    .name = "typing.TypeAliasType",
-    .basicsize = sizeof(typealiasobject),
-    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE | Py_TPFLAGS_HAVE_GC,
-    .slots = typealias_slots,
+PyTypeObject _PyTypeAlias_Type = {
+    PyVarObject_HEAD_INIT(&PyType_Type, 0)
+    .tp_name = "typing.TypeAliasType",
+    .tp_basicsize = sizeof(typealiasobject),
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE | Py_TPFLAGS_HAVE_GC,
+    .tp_doc = typealias_doc,
+    .tp_members = typealias_members,
+    .tp_methods = typealias_methods,
+    .tp_getset = typealias_getset,
+    .tp_alloc = PyType_GenericAlloc,
+    .tp_dealloc = typealias_dealloc,
+    .tp_new = typealias_new,
+    .tp_free = PyObject_GC_Del,
+    .tp_traverse = (traverseproc)typealias_traverse,
+    .tp_clear = (inquiry)typealias_clear,
+    .tp_repr = typealias_repr,
+    .tp_as_number = &typealias_as_number,
+    .tp_as_mapping = &typealias_as_mapping,
 };
 
 PyObject *
@@ -1445,7 +1475,8 @@ _Py_make_typealias(PyThreadState* unused, PyObject *args)
     assert(PyUnicode_Check(name));
     PyObject *type_params = PyTuple_GET_ITEM(args, 1);
     PyObject *compute_value = PyTuple_GET_ITEM(args, 2);
-    return (PyObject *)typealias_alloc(name, type_params, compute_value, NULL);
+    assert(PyFunction_Check(compute_value));
+    return (PyObject *)typealias_alloc(name, type_params, compute_value, NULL, NULL);
 }
 
 PyDoc_STRVAR(generic_doc,
@@ -1603,7 +1634,6 @@ int _Py_initialize_generic(PyInterpreterState *interp)
     MAKE_TYPE(paramspec);
     MAKE_TYPE(paramspecargs);
     MAKE_TYPE(paramspeckwargs);
-    MAKE_TYPE(typealias);
 #undef MAKE_TYPE
     return 0;
 }
@@ -1616,5 +1646,4 @@ void _Py_clear_generic_types(PyInterpreterState *interp)
     Py_CLEAR(interp->cached_objects.paramspec_type);
     Py_CLEAR(interp->cached_objects.paramspecargs_type);
     Py_CLEAR(interp->cached_objects.paramspeckwargs_type);
-    Py_CLEAR(interp->cached_objects.typealias_type);
 }

Objects/unionobject.c

@@ -1,6 +1,7 @@
 // types.UnionType -- used to represent e.g. Union[int, str], int | str
 #include "Python.h"
 #include "pycore_object.h"  // _PyObject_GC_TRACK/UNTRACK
+#include "pycore_typevarobject.h"  // _PyTypeAlias_Type
 #include "pycore_unionobject.h"
 #include "structmember.h"
 
@@ -150,11 +151,11 @@ is_unionable(PyObject *obj)
     if (obj == Py_None ||
         PyType_Check(obj) ||
         _PyGenericAlias_Check(obj) ||
-        _PyUnion_Check(obj)) {
+        _PyUnion_Check(obj) ||
+        Py_IS_TYPE(obj, &_PyTypeAlias_Type)) {
         return 1;
     }
-    PyInterpreterState *interp = PyInterpreterState_Get();
-    return Py_IS_TYPE(obj, interp->cached_objects.typealias_type);
+    return 0;
 }
 
 PyObject *

Tools/c-analyzer/cpython/globals-to-fix.tsv

@@ -90,6 +90,7 @@ Objects/typeobject.c	-	_PyBufferWrapper_Type	-
 Objects/typeobject.c	-	PyBaseObject_Type	-
 Objects/typeobject.c	-	PySuper_Type	-
 Objects/typeobject.c	-	PyType_Type	-
+Objects/typevarobject.c	-	_PyTypeAlias_Type	-
 Objects/unicodeobject.c	-	PyUnicodeIter_Type	-
 Objects/unicodeobject.c	-	PyUnicode_Type	-
 Objects/weakrefobject.c	-	_PyWeakref_CallableProxyType	-