From b81ca14beccf2ced29905fd54dd0ce7ec93978a4 Mon Sep 17 00:00:00 2001
From: herquan <31046219+herquan@users.noreply.github.com>
Date: Mon, 16 May 2022 19:50:51 -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 +
 guide/src/building_and_distribution.md |  4 ++++
 pyo3-ffi/src/import.rs                 |  2 +-
 src/macros.rs                          | 22 ++++++++++++++++++
 tests/test_append_to_inittab.rs        | 32 ++++++++++++++++++++++++++
 5 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_append_to_inittab.rs

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1a184d25e81..7cdf41919ec 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Implement `ToPyObject` for `[T; N]`. [#2313](https://github.com/PyO3/pyo3/pull/2313)
 - Added the internal `IntoPyResult` trait to give better error messages when function return types do not implement `IntoPy`. [#2326](https://github.com/PyO3/pyo3/pull/2326)
 - 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 macro `append_to_inittab`. [#2377](https://github.com/PyO3/pyo3/pull/2377)
 
 ### Changed
 
diff --git a/guide/src/building_and_distribution.md b/guide/src/building_and_distribution.md
index 60332c9ac6d..102d592d040 100644
--- a/guide/src/building_and_distribution.md
+++ b/guide/src/building_and_distribution.md
@@ -225,6 +225,10 @@ The known complications are:
 
 If you encounter these or other complications when linking the interpreter statically, discuss them on [issue 416 on PyO3's Github](https://github.com/PyO3/pyo3/issues/416). It is hoped that eventually that discussion will contain enough information and solutions that PyO3 can offer first-class support for static embedding.
 
+### Import your module when embedding the CPython interpreter
+
+When you run your Rust binary with CPython (not PyPy) interpreter, your newly created module won't be initialized unless the function defined with macro `#[pymodule]` is added to a table called `PyImport_Inittab`. This means Python statements like `import your_new_module` run by your Rust binary will fail. You can use macro `append_to_inittab` before function `prepare_freethreaded_python` being called to add the module function into that table. Also [`auto-initialize`](features.md#auto-initialize) needs to be turned off in such scenario.
+
 ## Cross Compiling
 
 Thanks to Rust's great cross-compilation support, cross-compiling using PyO3 is relatively straightforward. To get started, you'll need a few pieces of software:
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();
+    })
+}