bpo-1635741: Port _dbm module to multiphase initialization (GH-20848)

1 files changed, 170 insertions, 132 deletions

diff --git a/Modules/_dbmmodule.c b/Modules/_dbmmodule.c
index 80a0503..97772a0 100644
--- a/Modules/_dbmmodule.c
+++ b/Modules/_dbmmodule.c
@@ -28,6 +28,19 @@ static const char which_dbm[] = "Berkeley DB";
 #error "No ndbm.h available!"
 #endif
 
+typedef struct {
+    PyTypeObject *dbm_type;
+    PyObject *dbm_error;
+} _dbm_state;
+
+static inline _dbm_state*
+get_dbm_state(PyObject *module)
+{
+    void *state = PyModule_GetState(module);
+    assert(state != NULL);
+    return (_dbm_state *)state;
+}
+
 /*[clinic input]
 module _dbm
 class _dbm.dbm "dbmobject *" "&Dbmtype"
@@ -43,28 +56,25 @@ typedef struct {
 
 #include "clinic/_dbmmodule.c.h"
 
-static PyTypeObject Dbmtype;
-
-#define is_dbmobject(v) Py_IS_TYPE(v, &Dbmtype)
-#define check_dbmobject_open(v) if ((v)->di_dbm == NULL) \
-               { PyErr_SetString(DbmError, "DBM object has already been closed"); \
-                 return NULL; }
-
-static PyObject *DbmError;
+#define check_dbmobject_open(v, err)                                \
+    if ((v)->di_dbm == NULL) {                                      \
+        PyErr_SetString(err, "DBM object has already been closed"); \
+        return NULL;                                                \
+    }
 
 static PyObject *
-newdbmobject(const char *file, int flags, int mode)
+newdbmobject(_dbm_state *state, const char *file, int flags, int mode)
 {
     dbmobject *dp;
 
-    dp = PyObject_New(dbmobject, &Dbmtype);
+    dp = PyObject_New(dbmobject, state->dbm_type);
     if (dp == NULL)
         return NULL;
     dp->di_size = -1;
     dp->flags = flags;
     /* See issue #19296 */
     if ( (dp->di_dbm = dbm_open((char *)file, flags, mode)) == 0 ) {
-        PyErr_SetFromErrnoWithFilename(DbmError, file);
+        PyErr_SetFromErrnoWithFilename(state->dbm_error, file);
         Py_DECREF(dp);
         return NULL;
     }
@@ -76,16 +86,21 @@ newdbmobject(const char *file, int flags, int mode)
 static void
 dbm_dealloc(dbmobject *dp)
 {
-    if ( dp->di_dbm )
+    if (dp->di_dbm) {
         dbm_close(dp->di_dbm);
-    PyObject_Del(dp);
+    }
+    PyTypeObject *tp = Py_TYPE(dp);
+    tp->tp_free(dp);
+    Py_DECREF(tp);
 }
 
 static Py_ssize_t
 dbm_length(dbmobject *dp)
 {
+    _dbm_state *state = PyType_GetModuleState(Py_TYPE(dp));
+    assert(state != NULL);
     if (dp->di_dbm == NULL) {
-             PyErr_SetString(DbmError, "DBM object has already been closed");
+             PyErr_SetString(state->dbm_error, "DBM object has already been closed");
              return -1;
     }
     if ( dp->di_size < 0 ) {
@@ -106,12 +121,14 @@ dbm_subscript(dbmobject *dp, PyObject *key)
 {
     datum drec, krec;
     Py_ssize_t tmp_size;
-
-    if (!PyArg_Parse(key, "s#", &krec.dptr, &tmp_size) )
+    _dbm_state *state = PyType_GetModuleState(Py_TYPE(dp));
+    assert(state != NULL);
+    if (!PyArg_Parse(key, "s#", &krec.dptr, &tmp_size)) {
         return NULL;
+    }
 
     krec.dsize = tmp_size;
-    check_dbmobject_open(dp);
+    check_dbmobject_open(dp, state->dbm_error);
     drec = dbm_fetch(dp->di_dbm, krec);
     if ( drec.dptr == 0 ) {
         PyErr_SetObject(PyExc_KeyError, key);
@@ -119,7 +136,7 @@ dbm_subscript(dbmobject *dp, PyObject *key)
     }
     if ( dbm_error(dp->di_dbm) ) {
         dbm_clearerr(dp->di_dbm);
-        PyErr_SetString(DbmError, "");
+        PyErr_SetString(state->dbm_error, "");
         return NULL;
     }
     return PyBytes_FromStringAndSize(drec.dptr, drec.dsize);
@@ -136,9 +153,11 @@ dbm_ass_sub(dbmobject *dp, PyObject *v, PyObject *w)
                         "dbm mappings have bytes or string keys only");
         return -1;
     }
+    _dbm_state *state = PyType_GetModuleState(Py_TYPE(dp));
+    assert(state != NULL);
     krec.dsize = tmp_size;
     if (dp->di_dbm == NULL) {
-             PyErr_SetString(DbmError, "DBM object has already been closed");
+             PyErr_SetString(state->dbm_error, "DBM object has already been closed");
              return -1;
     }
     dp->di_size = -1;
@@ -151,7 +170,7 @@ dbm_ass_sub(dbmobject *dp, PyObject *v, PyObject *w)
                 PyErr_SetObject(PyExc_KeyError, v);
             }
             else {
-                PyErr_SetString(DbmError, "cannot delete item from database");
+                PyErr_SetString(state->dbm_error, "cannot delete item from database");
             }
             return -1;
         }
@@ -164,25 +183,19 @@ dbm_ass_sub(dbmobject *dp, PyObject *v, PyObject *w)
         drec.dsize = tmp_size;
         if ( dbm_store(dp->di_dbm, krec, drec, DBM_REPLACE) < 0 ) {
             dbm_clearerr(dp->di_dbm);
-            PyErr_SetString(DbmError,
+            PyErr_SetString(state->dbm_error,
                             "cannot add item to database");
             return -1;
         }
     }
     if ( dbm_error(dp->di_dbm) ) {
         dbm_clearerr(dp->di_dbm);
-        PyErr_SetString(DbmError, "");
+        PyErr_SetString(state->dbm_error, "");
         return -1;
     }
     return 0;
 }
 
-static PyMappingMethods dbm_as_mapping = {
-    (lenfunc)dbm_length,                /*mp_length*/
-    (binaryfunc)dbm_subscript,          /*mp_subscript*/
-    (objobjargproc)dbm_ass_sub,         /*mp_ass_subscript*/
-};
-
 /*[clinic input]
 _dbm.dbm.close
 
@@ -193,8 +206,9 @@ static PyObject *
 _dbm_dbm_close_impl(dbmobject *self)
 /*[clinic end generated code: output=c8dc5b6709600b86 input=046db72377d51be8]*/
 {
-    if (self->di_dbm)
+    if (self->di_dbm) {
         dbm_close(self->di_dbm);
+    }
     self->di_dbm = NULL;
     Py_RETURN_NONE;
 }
@@ -202,21 +216,26 @@ _dbm_dbm_close_impl(dbmobject *self)
 /*[clinic input]
 _dbm.dbm.keys
 
+    cls: defining_class
+
 Return a list of all keys in the database.
 [clinic start generated code]*/
 
 static PyObject *
-_dbm_dbm_keys_impl(dbmobject *self)
-/*[clinic end generated code: output=434549f7c121b33c input=d210ba778cd9c68a]*/
+_dbm_dbm_keys_impl(dbmobject *self, PyTypeObject *cls)
+/*[clinic end generated code: output=f2a593b3038e5996 input=d3706a28fc051097]*/
 {
     PyObject *v, *item;
     datum key;
     int err;
 
-    check_dbmobject_open(self);
+    _dbm_state *state = PyType_GetModuleState(cls);
+    assert(state != NULL);
+    check_dbmobject_open(self, state->dbm_error);
     v = PyList_New(0);
-    if (v == NULL)
+    if (v == NULL) {
         return NULL;
+    }
     for (key = dbm_firstkey(self->di_dbm); key.dptr;
          key = dbm_nextkey(self->di_dbm)) {
         item = PyBytes_FromStringAndSize(key.dptr, key.dsize);
@@ -241,8 +260,10 @@ dbm_contains(PyObject *self, PyObject *arg)
     datum key, val;
     Py_ssize_t size;
 
+    _dbm_state *state = PyType_GetModuleState(Py_TYPE(dp));
+    assert(state != NULL);
     if ((dp)->di_dbm == NULL) {
-        PyErr_SetString(DbmError,
+        PyErr_SetString(state->dbm_error,
                         "DBM object has already been closed");
          return -1;
     }
@@ -266,22 +287,9 @@ dbm_contains(PyObject *self, PyObject *arg)
     return val.dptr != NULL;
 }
 
-static PySequenceMethods dbm_as_sequence = {
-    0,                          /* sq_length */
-    0,                          /* sq_concat */
-    0,                          /* sq_repeat */
-    0,                          /* sq_item */
-    0,                          /* sq_slice */
-    0,                          /* sq_ass_item */
-    0,                          /* sq_ass_slice */
-    dbm_contains,               /* sq_contains */
-    0,                          /* sq_inplace_concat */
-    0,                          /* sq_inplace_repeat */
-};
-
 /*[clinic input]
 _dbm.dbm.get
-
+    cls: defining_class
     key: str(accept={str, robuffer}, zeroes=True)
     default: object = None
     /
@@ -290,19 +298,20 @@ Return the value for key if present, otherwise default.
 [clinic start generated code]*/
 
 static PyObject *
-_dbm_dbm_get_impl(dbmobject *self, const char *key,
+_dbm_dbm_get_impl(dbmobject *self, PyTypeObject *cls, const char *key,
                   Py_ssize_clean_t key_length, PyObject *default_value)
-/*[clinic end generated code: output=b44f95eba8203d93 input=b788eba0ffad2e91]*/
-/*[clinic end generated code: output=4f5c0e523eaf1251 input=9402c0af8582dc69]*/
+/*[clinic end generated code: output=34851b5dc1c664dc input=66b993b8349fa8c1]*/
 {
     datum dbm_key, val;
-
+    _dbm_state *state = PyType_GetModuleState(cls);
+    assert(state != NULL);
     dbm_key.dptr = (char *)key;
     dbm_key.dsize = key_length;
-    check_dbmobject_open(self);
+    check_dbmobject_open(self, state->dbm_error);
     val = dbm_fetch(self->di_dbm, dbm_key);
-    if (val.dptr != NULL)
+    if (val.dptr != NULL) {
         return PyBytes_FromStringAndSize(val.dptr, val.dsize);
+    }
 
     Py_INCREF(default_value);
     return default_value;
@@ -310,6 +319,7 @@ _dbm_dbm_get_impl(dbmobject *self, const char *key,
 
 /*[clinic input]
 _dbm.dbm.setdefault
+    cls: defining_class
     key: str(accept={str, robuffer}, zeroes=True)
     default: object(c_default="NULL") = b''
     /
@@ -320,24 +330,27 @@ If key is not in the database, it is inserted with default as the value.
 [clinic start generated code]*/
 
 static PyObject *
-_dbm_dbm_setdefault_impl(dbmobject *self, const char *key,
+_dbm_dbm_setdefault_impl(dbmobject *self, PyTypeObject *cls, const char *key,
                          Py_ssize_clean_t key_length,
                          PyObject *default_value)
-/*[clinic end generated code: output=52545886cf272161 input=bf40c48edaca01d6]*/
+/*[clinic end generated code: output=d5c68fe673886767 input=126a3ff15c5f8232]*/
 {
     datum dbm_key, val;
     Py_ssize_t tmp_size;
-
+    _dbm_state *state = PyType_GetModuleState(cls);
+    assert(state != NULL);
     dbm_key.dptr = (char *)key;
     dbm_key.dsize = key_length;
-    check_dbmobject_open(self);
+    check_dbmobject_open(self, state->dbm_error);
     val = dbm_fetch(self->di_dbm, dbm_key);
-    if (val.dptr != NULL)
+    if (val.dptr != NULL) {
         return PyBytes_FromStringAndSize(val.dptr, val.dsize);
+    }
     if (default_value == NULL) {
         default_value = PyBytes_FromStringAndSize(NULL, 0);
-        if (default_value == NULL)
+        if (default_value == NULL) {
             return NULL;
+        }
         val.dptr = NULL;
         val.dsize = 0;
     }
@@ -352,7 +365,7 @@ _dbm_dbm_setdefault_impl(dbmobject *self, const char *key,
     }
     if (dbm_store(self->di_dbm, dbm_key, val, DBM_INSERT) < 0) {
         dbm_clearerr(self->di_dbm);
-        PyErr_SetString(DbmError, "cannot add item to database");
+        PyErr_SetString(state->dbm_error, "cannot add item to database");
         Py_DECREF(default_value);
         return NULL;
     }
@@ -373,7 +386,6 @@ dbm__exit__(PyObject *self, PyObject *args)
     return _PyObject_CallMethodIdNoArgs(self, &PyId_close);
 }
 
-
 static PyMethodDef dbm_methods[] = {
     _DBM_DBM_CLOSE_METHODDEF
     _DBM_DBM_KEYS_METHODDEF
@@ -381,38 +393,29 @@ static PyMethodDef dbm_methods[] = {
     _DBM_DBM_SETDEFAULT_METHODDEF
     {"__enter__", dbm__enter__, METH_NOARGS, NULL},
     {"__exit__",  dbm__exit__, METH_VARARGS, NULL},
-    {NULL,              NULL}           /* sentinel */
+    {NULL,  NULL}           /* sentinel */
 };
 
-static PyTypeObject Dbmtype = {
-    PyVarObject_HEAD_INIT(NULL, 0)
-    "_dbm.dbm",
-    sizeof(dbmobject),
-    0,
-    (destructor)dbm_dealloc,  /*tp_dealloc*/
-    0,                            /*tp_vectorcall_offset*/
-    0,                        /*tp_getattr*/
-    0,                            /*tp_setattr*/
-    0,                            /*tp_as_async*/
-    0,                            /*tp_repr*/
-    0,                            /*tp_as_number*/
-    &dbm_as_sequence,             /*tp_as_sequence*/
-    &dbm_as_mapping,              /*tp_as_mapping*/
-    0,                    /*tp_hash*/
-    0,                    /*tp_call*/
-    0,                    /*tp_str*/
-    0,                    /*tp_getattro*/
-    0,                    /*tp_setattro*/
-    0,                    /*tp_as_buffer*/
-    Py_TPFLAGS_DEFAULT,   /*tp_flags*/
-    0,                        /*tp_doc*/
-    0,                        /*tp_traverse*/
-    0,                        /*tp_clear*/
-    0,                        /*tp_richcompare*/
-    0,                        /*tp_weaklistoffset*/
-    0,                        /*tp_iter*/
-    0,                        /*tp_iternext*/
-    dbm_methods,          /*tp_methods*/
+static PyType_Slot dbmtype_spec_slots[] = {
+    {Py_tp_dealloc, dbm_dealloc},
+    {Py_tp_methods, dbm_methods},
+    {Py_sq_contains, dbm_contains},
+    {Py_mp_length, dbm_length},
+    {Py_mp_subscript, dbm_subscript},
+    {Py_mp_ass_subscript, dbm_ass_sub},
+    {0, 0}
+};
+
+
+static PyType_Spec dbmtype_spec = {
+    .name = "_dbm.dbm",
+    .basicsize = sizeof(dbmobject),
+    // Calling PyType_GetModuleState() on a subclass is not safe.
+    // dbmtype_spec does not have Py_TPFLAGS_BASETYPE flag
+    // which prevents to create a subclass.
+    // So calling PyType_GetModuleState() in this file is always safe.
+    .flags = Py_TPFLAGS_DEFAULT,
+    .slots = dbmtype_spec_slots,
 };
 
 /* ----------------------------------------------------------------- */
@@ -443,19 +446,26 @@ dbmopen_impl(PyObject *module, PyObject *filename, const char *flags,
 /*[clinic end generated code: output=9527750f5df90764 input=376a9d903a50df59]*/
 {
     int iflags;
-
-    if ( strcmp(flags, "r") == 0 )
+    _dbm_state *state =  get_dbm_state(module);
+    assert(state != NULL);
+    if (strcmp(flags, "r") == 0) {
         iflags = O_RDONLY;
-    else if ( strcmp(flags, "w") == 0 )
+    }
+    else if (strcmp(flags, "w") == 0) {
         iflags = O_RDWR;
-    else if ( strcmp(flags, "rw") == 0 ) /* B/W compat */
+    }
+    else if (strcmp(flags, "rw") == 0) {
+        /* Backward compatibility */
         iflags = O_RDWR|O_CREAT;
-    else if ( strcmp(flags, "c") == 0 )
+    }
+    else if (strcmp(flags, "c") == 0) {
         iflags = O_RDWR|O_CREAT;
-    else if ( strcmp(flags, "n") == 0 )
+    }
+    else if (strcmp(flags, "n") == 0) {
         iflags = O_RDWR|O_CREAT|O_TRUNC;
+    }
     else {
-        PyErr_SetString(DbmError,
+        PyErr_SetString(state->dbm_error,
                         "arg 2 to open should be 'r', 'w', 'c', or 'n'");
         return NULL;
     }
@@ -470,7 +480,7 @@ dbmopen_impl(PyObject *module, PyObject *filename, const char *flags,
         PyErr_SetString(PyExc_ValueError, "embedded null character");
         return NULL;
     }
-    PyObject *self = newdbmobject(name, iflags, mode);
+    PyObject *self = newdbmobject(state, name, iflags, mode);
     Py_DECREF(filenamebytes);
     return self;
 }
@@ -480,42 +490,70 @@ static PyMethodDef dbmmodule_methods[] = {
     { 0, 0 },
 };
 
+static int
+_dbm_exec(PyObject *module)
+{
+    _dbm_state *state = get_dbm_state(module);
+    state->dbm_type = (PyTypeObject *)PyType_FromModuleAndSpec(module,
+                                                        &dbmtype_spec, NULL);
+    if (state->dbm_type == NULL) {
+        return -1;
+    }
+    state->dbm_error = PyErr_NewException("_dbm.error", PyExc_OSError, NULL);
+    if (state->dbm_error == NULL) {
+        return -1;
+    }
+    if (PyModule_AddStringConstant(module, "library", which_dbm) < 0) {
+        return -1;
+    }
+    if (PyModule_AddType(module, (PyTypeObject *)state->dbm_error) < 0) {
+        return -1;
+    }
+    return 0;
+}
+
+static int
+_dbm_module_traverse(PyObject *module, visitproc visit, void *arg)
+{
+    _dbm_state *state = get_dbm_state(module);
+    Py_VISIT(state->dbm_error);
+    Py_VISIT(state->dbm_type);
+    return 0;
+}
+
+static int
+_dbm_module_clear(PyObject *module)
+{
+    _dbm_state *state = get_dbm_state(module);
+    Py_CLEAR(state->dbm_error);
+    Py_CLEAR(state->dbm_type);
+    return 0;
+}
+
+static void
+_dbm_module_free(void *module)
+{
+    _dbm_module_clear((PyObject *)module);
+}
+
+static PyModuleDef_Slot _dbmmodule_slots[] = {
+    {Py_mod_exec, _dbm_exec},
+    {0, NULL}
+};
 
 static struct PyModuleDef _dbmmodule = {
     PyModuleDef_HEAD_INIT,
-    "_dbm",
-    NULL,
-    -1,
-    dbmmodule_methods,
-    NULL,
-    NULL,
-    NULL,
-    NULL
+    .m_name = "_dbm",
+    .m_size = sizeof(_dbm_state),
+    .m_methods = dbmmodule_methods,
+    .m_slots = _dbmmodule_slots,
+    .m_traverse = _dbm_module_traverse,
+    .m_clear = _dbm_module_clear,
+    .m_free = _dbm_module_free,
 };
 
 PyMODINIT_FUNC
-PyInit__dbm(void) {
-    PyObject *m, *d, *s;
-
-    if (PyType_Ready(&Dbmtype) < 0)
-        return NULL;
-    m = PyModule_Create(&_dbmmodule);
-    if (m == NULL)
-        return NULL;
-    d = PyModule_GetDict(m);
-    if (DbmError == NULL)
-        DbmError = PyErr_NewException("_dbm.error",
-                                      PyExc_OSError, NULL);
-    s = PyUnicode_FromString(which_dbm);
-    if (s != NULL) {
-        PyDict_SetItemString(d, "library", s);
-        Py_DECREF(s);
-    }
-    if (DbmError != NULL)
-        PyDict_SetItemString(d, "error", DbmError);
-    if (PyErr_Occurred()) {
-        Py_DECREF(m);
-        m = NULL;
-    }
-    return m;
+PyInit__dbm(void)
+{
+    return PyModuleDef_Init(&_dbmmodule);
 }