From f9644e095c1cd4642509f364ae2bb6730f29e0ef 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
---
 CHANGELOG.md                    |  1 +
 pyo3-ffi/src/import.rs          |  2 +-
 src/macros.rs                   | 22 ++++++++++++++++++++++
 tests/test_append_to_inittab.rs | 32 ++++++++++++++++++++++++++++++++
 4 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_append_to_inittab.rs

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 42a6a68401c..11e2762fbff 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add `PyDictKeys`, `PyDictValues` and `PyDictItems` Rust types to represent `dict_keys`, `dict_values` and `dict_items` types. [#2358](https://github.com/PyO3/pyo3/pull/2358)
 - Add an experimental `generate-import-lib` feature to support auto-generating non-abi3 python import libraries for Windows targets. [#2364](https://github.com/PyO3/pyo3/pull/2364)
 - Add FFI definition `Py_ExitStatusException`. [#2374](https://github.com/PyO3/pyo3/pull/2374)
+- Add macro `append_to_inittab`. [#2377](https://github.com/PyO3/pyo3/pull/2377)
 
 ### Changed
 
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..90bfa266e6a 100644
--- a/src/macros.rs
+++ b/src/macros.rs
@@ -150,3 +150,25 @@ macro_rules! wrap_pymodule {
         }
     };
 }
+
+#[cfg(not(PyPy))]
+/// 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_append_to_inittab.rs b/tests/test_append_to_inittab.rs
new file mode 100644
index 00000000000..87f4cf02789
--- /dev/null
+++ b/tests/test_append_to_inittab.rs
@@ -0,0 +1,32 @@
+#![cfg(feature = "macros")]
+use pyo3::prelude::*;
+
+#[pyfunction]
+fn foo() -> usize {
+    123
+}
+
+#[pymodule]
+fn module_with_functions(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(foo, m)?).unwrap();
+    Ok(())
+}
+
+#[cfg(not(PyPy))]
+#[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.foo() == 123
+"#,
+            None,
+            None,
+        )
+        .map_err(|e| e.print(py))
+        .unwrap();
+    })
+}