bpo-44553 : Implement GC methods for types.Union (GH-26993)

3 files changed, 38 insertions, 5 deletions

diff --git a/Lib/test/test_types.py b/Lib/test/test_types.py
index d3e315a..ae7b17b 100644
--- a/Lib/test/test_types.py
+++ b/Lib/test/test_types.py
@@ -1,8 +1,9 @@
 # Python test set -- part 6, built-in types
 
-from test.support import run_with_locale
+from test.support import run_with_locale, cpython_only
 import collections.abc
 from collections import namedtuple
+import gc
 import inspect
 import pickle
 import locale
@@ -756,6 +757,23 @@ class TypesTests(unittest.TestCase):
         with self.assertRaises(ZeroDivisionError):
             str | TypeVar()
 
+    @cpython_only
+    def test_or_type_operator_reference_cycle(self):
+        if not hasattr(sys, 'gettotalrefcount'):
+            self.skipTest('Cannot get total reference count.')
+        gc.collect()
+        before = sys.gettotalrefcount()
+        for _ in range(30):
+            T = typing.TypeVar('T')
+            U = int | list[T]
+            T.blah = U
+            del T
+            del U
+        gc.collect()
+        leeway = 15
+        self.assertLessEqual(sys.gettotalrefcount() - before, leeway,
+                             msg='Check for union reference leak.')
+
     def test_ellipsis_type(self):
         self.assertIsInstance(Ellipsis, types.EllipsisType)
 
diff --git a/Misc/NEWS.d/next/Core and Builtins/2021-07-02-22-54-41.bpo-44553.l9YqGg.rst b/Misc/NEWS.d/next/Core and Builtins/2021-07-02-22-54-41.bpo-44553.l9YqGg.rst
new file mode 100644
index 0000000..e97df02
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2021-07-02-22-54-41.bpo-44553.l9YqGg.rst
@@ -0,0 +1,2 @@
+Implement GC methods for ``types.Union`` to break reference cycles and
+prevent memory leaks.
diff --git a/Objects/unionobject.c b/Objects/unionobject.c
index a66d615..cc7181d 100644
--- a/Objects/unionobject.c
+++ b/Objects/unionobject.c
@@ -1,5 +1,6 @@
 // types.Union -- used to represent e.g. Union[int, str], int | str
 #include "Python.h"
+#include "pycore_object.h"  // _PyObject_GC_TRACK/UNTRACK
 #include "pycore_unionobject.h"
 #include "structmember.h"
 
@@ -14,10 +15,20 @@ unionobject_dealloc(PyObject *self)
 {
     unionobject *alias = (unionobject *)self;
 
+    _PyObject_GC_UNTRACK(self);
+
     Py_XDECREF(alias->args);
     Py_TYPE(self)->tp_free(self);
 }
 
+static int
+union_traverse(PyObject *self, visitproc visit, void *arg)
+{
+    unionobject *alias = (unionobject *)self;
+    Py_VISIT(alias->args);
+    return 0;
+}
+
 static Py_hash_t
 union_hash(PyObject *self)
 {
@@ -437,8 +448,9 @@ PyTypeObject _Py_UnionType = {
     .tp_basicsize = sizeof(unionobject),
     .tp_dealloc = unionobject_dealloc,
     .tp_alloc = PyType_GenericAlloc,
-    .tp_free = PyObject_Del,
-    .tp_flags = Py_TPFLAGS_DEFAULT,
+    .tp_free = PyObject_GC_Del,
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
+    .tp_traverse = union_traverse,
     .tp_hash = union_hash,
     .tp_getattro = PyObject_GenericGetAttr,
     .tp_members = union_members,
@@ -472,15 +484,16 @@ _Py_Union(PyObject *args)
         }
     }
 
-    result = PyObject_New(unionobject, &_Py_UnionType);
+    result = PyObject_GC_New(unionobject, &_Py_UnionType);
     if (result == NULL) {
         return NULL;
     }
 
     result->args = dedup_and_flatten_args(args);
     if (result->args == NULL) {
-        Py_DECREF(result);
+        PyObject_GC_Del(result);
         return NULL;
     }
+    _PyObject_GC_TRACK(result);
     return (PyObject*)result;
 }