gh-103583: Add ref. dependency between multibytecodec modules (#103589)

erlend-aasland · web-flow · commit a6b07b5a345f · 2023-04-19T10:02:17.000-06:00
diff --git a/Modules/cjkcodecs/cjkcodecs.h b/Modules/cjkcodecs/cjkcodecs.h
@@ -284,18 +284,45 @@ getmultibytecodec(void)
     return _PyImport_GetModuleAttrString("_multibytecodec", "__create_codec");
 }
 
+static void
+destroy_codec_capsule(PyObject *capsule)
+{
+    void *ptr = PyCapsule_GetPointer(capsule, CODEC_CAPSULE);
+    codec_capsule *data = (codec_capsule *)ptr;
+    Py_DECREF(data->cjk_module);
+    PyMem_Free(ptr);
+}
+
+static codec_capsule *
+capsulate_codec(PyObject *mod, const MultibyteCodec *codec)
+{
+    codec_capsule *data = PyMem_Malloc(sizeof(codec_capsule));
+    if (data == NULL) {
+        PyErr_NoMemory();
+        return NULL;
+    }
+    data->codec = codec;
+    data->cjk_module = Py_NewRef(mod);
+    return data;
+}
+
 static PyObject *
-_getcodec(const MultibyteCodec *codec)
+_getcodec(PyObject *self, const MultibyteCodec *codec)
 {
     PyObject *cofunc = getmultibytecodec();
     if (cofunc == NULL) {
         return NULL;
     }
 
-    PyObject *codecobj = PyCapsule_New((void *)codec,
-                                       PyMultibyteCodec_CAPSULE_NAME,
-                                       NULL);
+    codec_capsule *data = capsulate_codec(self, codec);
+    if (data == NULL) {
+        Py_DECREF(cofunc);
+        return NULL;
+    }
+    PyObject *codecobj = PyCapsule_New(data, CODEC_CAPSULE,
+                                       destroy_codec_capsule);
     if (codecobj == NULL) {
+        PyMem_Free(data);
         Py_DECREF(cofunc);
         return NULL;
     }
@@ -323,7 +350,7 @@ getcodec(PyObject *self, PyObject *encoding)
     for (int i = 0; i < st->num_codecs; i++) {
         const MultibyteCodec *codec = &st->codec_list[i];
         if (strcmp(codec->encoding, enc) == 0) {
-            return _getcodec(codec);
+            return _getcodec(self, codec);
         }
     }
 
@@ -352,8 +379,7 @@ register_maps(PyObject *module)
         char mhname[256] = "__map_";
         strcpy(mhname + sizeof("__map_") - 1, h->charset);
 
-        PyObject *capsule = PyCapsule_New((void *)h,
-                                          PyMultibyteCodec_CAPSULE_NAME, NULL);
+        PyObject *capsule = PyCapsule_New((void *)h, MAP_CAPSULE, NULL);
         if (capsule == NULL) {
             return -1;
         }
@@ -417,14 +443,14 @@ importmap(const char *modname, const char *symbol,
     o = PyObject_GetAttrString(mod, symbol);
     if (o == NULL)
         goto errorexit;
-    else if (!PyCapsule_IsValid(o, PyMultibyteCodec_CAPSULE_NAME)) {
+    else if (!PyCapsule_IsValid(o, MAP_CAPSULE)) {
         PyErr_SetString(PyExc_ValueError,
                         "map data must be a Capsule.");
         goto errorexit;
     }
     else {
         struct dbcs_map *map;
-        map = PyCapsule_GetPointer(o, PyMultibyteCodec_CAPSULE_NAME);
+        map = PyCapsule_GetPointer(o, MAP_CAPSULE);
         if (encmap != NULL)
             *encmap = map->encmap;
         if (decmap != NULL)
diff --git a/Modules/cjkcodecs/multibytecodec.c b/Modules/cjkcodecs/multibytecodec.c
@@ -720,9 +720,17 @@ static struct PyMethodDef multibytecodec_methods[] = {
 };
 
 static int
-multibytecodec_traverse(PyObject *self, visitproc visit, void *arg)
+multibytecodec_clear(MultibyteCodecObject *self)
+{
+    Py_CLEAR(self->cjk_module);
+    return 0;
+}
+
+static int
+multibytecodec_traverse(MultibyteCodecObject *self, visitproc visit, void *arg)
 {
     Py_VISIT(Py_TYPE(self));
+    Py_VISIT(self->cjk_module);
     return 0;
 }
 
@@ -731,6 +739,7 @@ multibytecodec_dealloc(MultibyteCodecObject *self)
 {
     PyObject_GC_UnTrack(self);
     PyTypeObject *tp = Py_TYPE(self);
+    (void)multibytecodec_clear(self);
     tp->tp_free(self);
     Py_DECREF(tp);
 }
@@ -740,6 +749,7 @@ static PyType_Slot multibytecodec_slots[] = {
     {Py_tp_getattro, PyObject_GenericGetAttr},
     {Py_tp_methods, multibytecodec_methods},
     {Py_tp_traverse, multibytecodec_traverse},
+    {Py_tp_clear, multibytecodec_clear},
     {0, NULL},
 };
 
@@ -1953,14 +1963,14 @@ _multibytecodec___create_codec(PyObject *module, PyObject *arg)
 /*[clinic end generated code: output=cfa3dce8260e809d input=6840b2a6b183fcfa]*/
 {
     MultibyteCodecObject *self;
-    const MultibyteCodec *codec;
 
-    if (!PyCapsule_IsValid(arg, PyMultibyteCodec_CAPSULE_NAME)) {
+    if (!PyCapsule_IsValid(arg, CODEC_CAPSULE)) {
         PyErr_SetString(PyExc_ValueError, "argument type invalid");
         return NULL;
     }
 
-    codec = PyCapsule_GetPointer(arg, PyMultibyteCodec_CAPSULE_NAME);
+    codec_capsule *data = PyCapsule_GetPointer(arg, CODEC_CAPSULE);
+    const MultibyteCodec *codec = data->codec;
     if (codec->codecinit != NULL && codec->codecinit(codec->config) != 0)
         return NULL;
 
@@ -1969,6 +1979,7 @@ _multibytecodec___create_codec(PyObject *module, PyObject *arg)
     if (self == NULL)
         return NULL;
     self->codec = codec;
+    self->cjk_module = Py_NewRef(data->cjk_module);
 
     PyObject_GC_Track(self);
     return (PyObject *)self;
diff --git a/Modules/cjkcodecs/multibytecodec.h b/Modules/cjkcodecs/multibytecodec.h
@@ -63,6 +63,7 @@ typedef struct {
 typedef struct {
     PyObject_HEAD
     const MultibyteCodec *codec;
+    PyObject *cjk_module;
 } MultibyteCodecObject;
 
 #define MultibyteCodec_Check(state, op) Py_IS_TYPE((op), state->multibytecodec_type)
@@ -130,7 +131,13 @@ typedef struct {
 #define MBENC_FLUSH             0x0001 /* encode all characters encodable */
 #define MBENC_MAX               MBENC_FLUSH
 
-#define PyMultibyteCodec_CAPSULE_NAME "multibytecodec.__map_*"
+typedef struct {
+    const MultibyteCodec *codec;
+    PyObject *cjk_module;
+} codec_capsule;
+
+#define MAP_CAPSULE "multibytecodec.map"
+#define CODEC_CAPSULE "multibytecodec.codec"
 
 
 #ifdef __cplusplus
