Add module state

Author: ZipFile

src/sonyflake_turbo/_sonyflake.c

@@ -1,13 +1,30 @@
 #include <Python.h>
 #include <pythread.h>
-#include <stdatomic.h>
 #include <stdbool.h>
 #include <stdint.h>
-#include <time.h>
 
+#include "module.h"
 #include "sonyflake.h"
 #include "machine_ids.h"
 
+static int sonyflake_module_traverse(PyObject *m, visitproc visit, void *arg) {
+	struct sonyflake_module_state *state = PyModule_GetState(m);
+	Py_VISIT(state->sonyflake_cls);
+	Py_VISIT(state->machine_id_lcg_cls);
+	return 0;
+}
+
+static int sonyflake_module_clear(PyObject *m) {
+	struct sonyflake_module_state *state = PyModule_GetState(m);
+	Py_CLEAR(state->sonyflake_cls);
+	Py_CLEAR(state->machine_id_lcg_cls);
+	return 0;
+}
+
+static void sonyflake_module_free(void *m) {
+	sonyflake_module_clear(m);
+}
+
 static int sonyflake_exec(PyObject *module);
 
 PyModuleDef_Slot sonyflake_slots[] = {
@@ -25,8 +42,11 @@ struct PyModuleDef sonyflake_module = {
 	PyModuleDef_HEAD_INIT,
 	.m_name = MODULE_NAME,
 	.m_doc = "",
-	.m_size = 0,
+	.m_size = sizeof(struct sonyflake_module_state),
 	.m_slots = sonyflake_slots,
+	.m_traverse = sonyflake_module_traverse,
+	.m_clear = sonyflake_module_clear,
+	.m_free = sonyflake_module_free,
 };
 
 PyMODINIT_FUNC
@@ -36,26 +56,29 @@ PyInit__sonyflake(void)
 }
 
 static int sonyflake_exec(PyObject *module) {
-	PyObject *sonyflake_cls, *machine_id_lcg_cls;
+	struct sonyflake_module_state *state = PyModule_GetState(module);
+
+	state->sonyflake_cls = NULL;
+	state->machine_id_lcg_cls = NULL;
 
-	sonyflake_cls = PyType_FromSpec(&sonyflake_type_spec);
+	state->sonyflake_cls = PyType_FromModuleAndSpec(module, &sonyflake_type_spec, NULL);
 
-	if (!sonyflake_cls) {
+	if (!state->sonyflake_cls) {
 		goto err;
 	}
 
-	if (PyModule_AddObject(module, "SonyFlake", sonyflake_cls) < 0) {
-		goto err_sf;
+	if (PyModule_AddObjectRef(module, "SonyFlake", state->sonyflake_cls) < 0) {
+		goto err;
 	}
 
-	machine_id_lcg_cls = PyType_FromSpec(&machine_id_lcg_spec);
+	state->machine_id_lcg_cls = PyType_FromModuleAndSpec(module, &machine_id_lcg_spec, NULL);
 
-	if (!machine_id_lcg_cls) {
-		goto err_lcg;
+	if (!state->machine_id_lcg_cls) {
+		goto err;
 	}
 
-	if (PyModule_AddObject(module, "MachineIDLCG", machine_id_lcg_cls) < 0) {
-		goto err_lcg;
+	if (PyModule_AddObjectRef(module, "MachineIDLCG", state->machine_id_lcg_cls) < 0) {
+		goto err;
 	}
 
 	PyModule_AddIntMacro(module, SONYFLAKE_EPOCH);
@@ -68,10 +91,9 @@ static int sonyflake_exec(PyObject *module) {
 
 	return 0;
 
-err_lcg:
-	Py_DECREF(machine_id_lcg_cls);
-err_sf:
-	Py_DECREF(sonyflake_cls);
 err:
+	Py_CLEAR(state->machine_id_lcg_cls);
+	Py_CLEAR(state->sonyflake_cls);
+
 	return -1;
 }

src/sonyflake_turbo/module.h

@@ -0,0 +1,15 @@
+#pragma once
+
+#include <Python.h>
+#include <assert.h>
+
+#define MODULE_NAME "sonyflake_turbo"
+#define MODULE(self) ((struct sonyflake_module_state *) PyType_GetModuleState(Py_TYPE(self)))
+
+struct sonyflake_module_state {
+	PyObject *sonyflake_cls;
+	PyObject *machine_id_lcg_cls;
+	PyObject *sleep_wrapper_cls;
+};
+
+extern struct PyModuleDef sonyflake_module;

src/sonyflake_turbo/sonyflake.h

@@ -3,8 +3,8 @@
 #include <Python.h>
 #include <stdint.h>
 #include <time.h>
+#include "module.h"
 
-#define MODULE_NAME "sonyflake_turbo"
 #define SONYFLAKE_EPOCH 1409529600 // 2014-09-01 00:00:00 UTC
 #define SONYFLAKE_SEQUENCE_BITS 8
 #define SONYFLAKE_SEQUENCE_MAX ((1 << SONYFLAKE_SEQUENCE_BITS) - 1)