From a699244bcc22cdfd8db3cf83444561a27594ac51 Mon Sep 17 00:00:00 2001
From: herquan <31046219+herquan@users.noreply.github.com>
Date: Sun, 15 May 2022 20:18:11 -0700
Subject: [PATCH] Add macro append_to_inittab

Sometimes we need to debug in a real environment with our module installed. `append_to_inittab` will be a wrapper for PyImport_AppendInittab (https://docs.python.org/3/c-api/import.html#c.PyImport_AppendInittab) and help us to do this
---
 pyo3-ffi/src/import.rs |  2 +-
 src/macros.rs          | 21 +++++++++++++++++++++
 tests/test_module.rs   | 18 ++++++++++++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/pyo3-ffi/src/import.rs b/pyo3-ffi/src/import.rs
index 794e0ee5480..e00843466e8 100644
--- a/pyo3-ffi/src/import.rs
+++ b/pyo3-ffi/src/import.rs
@@ -76,6 +76,6 @@ extern "C" {
 
     pub fn PyImport_AppendInittab(
         name: *const c_char,
-        initfunc: Option<extern "C" fn() -> *mut PyObject>,
+        initfunc: Option<unsafe extern "C" fn() -> *mut PyObject>,
     ) -> c_int;
 }
diff --git a/src/macros.rs b/src/macros.rs
index 0a164b350f5..2788c622449 100644
--- a/src/macros.rs
+++ b/src/macros.rs
@@ -150,3 +150,24 @@ macro_rules! wrap_pymodule {
         }
     };
 }
+
+/// Add current module to the existing table of built-in modules.
+///
+/// Use it before [`prepare_freethreaded_python`](crate::prepare_freethreaded_python) and
+/// leave feature `auto-initialize` off
+#[macro_export]
+macro_rules! append_to_inittab {
+    ($module:ident) => {
+        unsafe {
+            assert_eq!(
+                $crate::ffi::Py_IsInitialized(),
+                0,
+                "called `append_to_inittab_impl` but a Python interpreter is already running."
+            );
+            $crate::ffi::PyImport_AppendInittab(
+                concat!(stringify!($module), "\0").as_ptr() as *const std::os::raw::c_char,
+                Option::Some($module::init),
+            );
+        }
+    };
+}
diff --git a/tests/test_module.rs b/tests/test_module.rs
index 7364d7ec64c..ee7dcb548b3 100644
--- a/tests/test_module.rs
+++ b/tests/test_module.rs
@@ -443,3 +443,21 @@ fn test_module_doc_hidden() {
         py_assert!(py, m, "m.__doc__ == ''");
     })
 }
+
+#[test]
+fn test_module_append_to_inittab() {
+    use pyo3::append_to_inittab;
+    append_to_inittab!(module_with_functions);
+    Python::with_gil(|py| {
+        py.run(
+            r#"
+import module_with_functions
+assert module_with_functions.double(3)==6
+        "#,
+            None,
+            None,
+        )
+        .map_err(|e| e.print(py))
+        .unwrap();
+    })
+}