From a734b3387135777b4f7b0436d880a92a13f1de73 Mon Sep 17 00:00:00 2001
From: Heran Quan <31046219+herquan@users.noreply.github.com>
Date: Mon, 9 May 2022 14:44:15 -0700
Subject: [PATCH] [DO NOT MERGE] [#2361]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 +-
 pyo3-macros-backend/src/lib.rs    |  2 +-
 pyo3-macros-backend/src/module.rs | 17 +++++++++++++++++
 pyo3-macros/src/lib.rs            | 11 ++++++++++-
 src/lib.rs                        |  1 +
 tests/test_module.rs              | 12 ++++++++++++
 6 files changed, 42 insertions(+), 3 deletions(-)

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/pyo3-macros-backend/src/lib.rs b/pyo3-macros-backend/src/lib.rs
index 7ed8e504f6a..10b1519c20e 100644
--- a/pyo3-macros-backend/src/lib.rs
+++ b/pyo3-macros-backend/src/lib.rs
@@ -29,7 +29,7 @@ mod pyproto;
 mod wrap;
 
 pub use frompyobject::build_derive_from_pyobject;
-pub use module::{process_functions_in_module, pymodule_impl, PyModuleOptions};
+pub use module::{process_functions_in_module, pymodule_impl, PyModuleOptions, append_to_inittab_impl};
 pub use pyclass::{build_py_class, build_py_enum, PyClassArgs};
 pub use pyfunction::{build_py_function, PyFunctionOptions};
 pub use pyimpl::{build_py_methods, PyClassMethodsType};
diff --git a/pyo3-macros-backend/src/module.rs b/pyo3-macros-backend/src/module.rs
index 85a94b84186..d7b13acad1f 100644
--- a/pyo3-macros-backend/src/module.rs
+++ b/pyo3-macros-backend/src/module.rs
@@ -90,6 +90,23 @@ pub fn pymodule_impl(
     }
 }
 
+pub fn append_to_inittab_impl(fnname: &Ident) -> TokenStream {
+    let name = fnname.unraw();
+    let cb_name = Ident::new(&format!("PyInit_{}", name), Span::call_site());
+    quote! {
+        unsafe{
+            use pyo3::ffi::PyImport_AppendInittab;
+            assert_eq!(
+                pyo3::ffi::Py_IsInitialized(),
+                0,
+                "called `append_to_inittab_impl` but a Python interpreter is already running."
+            );
+            // TODO: How to pass in pyo3?
+            pyo3::ffi::PyImport_AppendInittab(concat!(stringify!(#name), "\0").as_ptr() as *const std::os::raw::c_char, Option::Some(#cb_name));
+        }
+    }
+}
+
 /// Finds and takes care of the #[pyfn(...)] in `#[pymodule]`
 pub fn process_functions_in_module(func: &mut syn::ItemFn) -> syn::Result<()> {
     let mut stmts: Vec<syn::Stmt> = Vec::new();
diff --git a/pyo3-macros/src/lib.rs b/pyo3-macros/src/lib.rs
index 84160af2b95..412812266dc 100644
--- a/pyo3-macros/src/lib.rs
+++ b/pyo3-macros/src/lib.rs
@@ -9,7 +9,7 @@ use proc_macro::TokenStream;
 use proc_macro2::TokenStream as TokenStream2;
 use pyo3_macros_backend::{
     build_derive_from_pyobject, build_py_class, build_py_enum, build_py_function, build_py_methods,
-    get_doc, process_functions_in_module, pymodule_impl, wrap_pyfunction_impl, wrap_pymodule_impl,
+    get_doc, process_functions_in_module, pymodule_impl, append_to_inittab_impl,wrap_pyfunction_impl, wrap_pymodule_impl,
     PyClassArgs, PyClassMethodsType, PyFunctionOptions, PyModuleOptions, WrapPyFunctionArgs,
 };
 use quote::quote;
@@ -213,6 +213,15 @@ pub fn wrap_pymodule(input: TokenStream) -> TokenStream {
     wrap_pymodule_impl(path).unwrap_or_compile_error().into()
 }
 
+/// Add current module to the existing table of built-in modules.
+///
+/// Use before `prepare_freethreaded_python()`
+#[proc_macro]
+pub fn append_to_inittab(input: TokenStream) -> TokenStream {
+    let fnname = parse_macro_input!(input as syn::Ident);
+    append_to_inittab_impl(&fnname).into()
+}
+
 fn pyclass_impl(
     attrs: TokenStream,
     mut ast: syn::ItemStruct,
diff --git a/src/lib.rs b/src/lib.rs
index 654a14dc4d5..1d22dc67fc4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -407,6 +407,7 @@ pub use pyo3_macros::pyproto;
 #[cfg(feature = "macros")]
 pub use pyo3_macros::{
     pyclass, pyfunction, pymethods, pymodule, wrap_pyfunction, wrap_pymodule, FromPyObject,
+    append_to_inittab
 };
 
 #[cfg(feature = "macros")]
diff --git a/tests/test_module.rs b/tests/test_module.rs
index 7364d7ec64c..59d9f1c4068 100644
--- a/tests/test_module.rs
+++ b/tests/test_module.rs
@@ -4,6 +4,7 @@ use pyo3::prelude::*;
 
 use pyo3::py_run;
 use pyo3::types::{IntoPyDict, PyDict, PyTuple};
+use pyo3_macros::append_to_inittab;
 mod common;
 
 #[pyclass]
@@ -443,3 +444,14 @@ fn test_module_doc_hidden() {
         py_assert!(py, m, "m.__doc__ == ''");
     })
 }
+
+#[test]
+fn test_module_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();
+    })
+}
\ No newline at end of file