Implement PEP 489 - Multi-phase extension module initialization

Cargo.toml

@@ -78,6 +78,9 @@ abi3-py38 = ["abi3-py39", "pyo3-build-config/abi3-py38", "pyo3-ffi/abi3-py38"]
 abi3-py39 = ["abi3-py310", "pyo3-build-config/abi3-py39", "pyo3-ffi/abi3-py39"]
 abi3-py310 = ["abi3", "pyo3-build-config/abi3-py310", "pyo3-ffi/abi3-py310"]
 
+# Enable PEP 489 - Multi-phase extension module initialization
+pep489 = ["pyo3-macros/pep489"]
+
 # Changes `Python::with_gil` and `Python::acquire_gil` to automatically initialize the
 # Python interpreter if needed.
 auto-initialize = []
@@ -98,6 +101,7 @@ full = [
     "indexmap",
     "eyre",
     "anyhow",
+    "pep489",
 ]
 
 [[bench]]

examples/decorator/Cargo.toml

@@ -10,4 +10,7 @@ crate-type = ["cdylib"]
 [dependencies]
 pyo3 = { path = "../../", features = ["extension-module"] }
 
+[features]
+pep489 = ["pyo3/pep489"]
+
 [workspace]

examples/decorator/noxfile.py

@@ -2,8 +2,11 @@
 
 
 @nox.session
-def python(session):
+@nox.parametrize("cargo_features", ['""', "pep489"])
+def python(session, cargo_features):
     session.install("-rrequirements-dev.txt")
     session.install("maturin")
-    session.run_always("maturin", "develop")
+    session.run_always(
+        "maturin", "develop", f"--cargo-extra-args=--features {cargo_features}"
+    )
     session.run("pytest")

examples/maturin-starter/Cargo.toml

@@ -10,4 +10,7 @@ crate-type = ["cdylib"]
 [dependencies]
 pyo3 = { path = "../../", features = ["extension-module"] }
 
+[features]
+pep489 = ["pyo3/pep489"]
+
 [workspace]

examples/maturin-starter/noxfile.py

@@ -2,8 +2,11 @@
 
 
 @nox.session
-def python(session):
+@nox.parametrize("cargo_features", ['""', "pep489"])
+def python(session, cargo_features):
     session.install("-rrequirements-dev.txt")
     session.install("maturin")
-    session.run_always("maturin", "develop")
+    session.run_always(
+        "maturin", "develop", f"--cargo-extra-args=--features {cargo_features}"
+    )
     session.run("pytest")

examples/setuptools-rust-starter/Cargo.toml

@@ -10,4 +10,7 @@ crate-type = ["cdylib"]
 [dependencies]
 pyo3 = { path = "../../", features = ["extension-module"] }
 
+[features]
+pep489 = ["pyo3/pep489"]
+
 [workspace]

examples/word-count/Cargo.toml

@@ -11,4 +11,7 @@ crate-type = ["cdylib"]
 pyo3 = { path = "../..", features = ["extension-module"] }
 rayon = "1.0.2"
 
+[features]
+pep489 = ["pyo3/pep489"]
+
 [workspace]

examples/word-count/noxfile.py

@@ -4,10 +4,13 @@
 
 
 @nox.session
-def test(session):
+@nox.parametrize("cargo_features", ['""', "pep489"])
+def test(session, cargo_features):
     session.install("-rrequirements-dev.txt")
     session.install("maturin")
-    session.run_always("maturin", "develop")
+    session.run_always(
+        "maturin", "develop", f"--cargo-extra-args=--features {cargo_features}"
+    )
     session.run("pytest")
 
 

pyo3-macros-backend/Cargo.toml

@@ -16,6 +16,7 @@ edition = "2018"
 [dependencies]
 quote = { version = "1", default-features = false }
 proc-macro2 = { version = "1", default-features = false }
+pyo3-build-config = { path = "../pyo3-build-config", version = "0.16.2", features = ["resolve-config"] }
 
 [dependencies.syn]
 version = "1.0.56"
@@ -25,3 +26,4 @@ features = ["derive", "parsing", "printing", "clone-impls", "full", "extra-trait
 [features]
 pyproto = []
 abi3 = []
+pep489 = []

pyo3-macros-backend/src/module.rs

@@ -71,10 +71,12 @@ pub fn pymodule_impl(
     let name = options.name.unwrap_or_else(|| fnname.unraw());
     let krate = get_pyo3_crate(&options.krate);
     let cb_name = Ident::new(&format!("PyInit_{}", name), Span::call_site());
+    let slots_name = Ident::new(&format!("{}_slots", name), Span::call_site());
+    let mod_exec_name = Ident::new(&format!("{}_exec", name), Span::call_site());
 
     let module_def_name = module_def_ident(fnname);
 
-    quote! {
+    let single_phase_module = quote! {
         #[no_mangle]
         #[allow(non_snake_case)]
         /// This autogenerated function is called by the python interpreter when importing
@@ -87,9 +89,69 @@ pub fn pymodule_impl(
         #visibility static #module_def_name: #krate::impl_::pymodule::ModuleDef = unsafe {
             #krate::impl_::pymodule::ModuleDef::new(concat!(stringify!(#name), "\0"), #doc, #krate::impl_::pymodule::ModuleInitializer(#fnname))
         };
+    };
+
+    let multi_phase_module = quote! {
+        #[no_mangle]
+        #[allow(non_snake_case)]
+        /// This autogenerated function is called by the python interpreter when importing
+        /// the module.
+        pub unsafe extern "C" fn #cb_name() -> *mut #krate::ffi::PyObject {
+            unsafe { #krate::ffi::PyModuleDef_Init(#module_def_name.module_def()) }
+        }
+
+        pub unsafe extern "C" fn #mod_exec_name(module: *mut #krate::ffi::PyObject) -> ::std::os::raw::c_int {
+            use #krate::conversion::FromPyPointer;
+
+            let pool = #krate::GILPool::new();
+            let py = pool.python();
+            let module = #krate::types::PyModule::from_borrowed_ptr(py, module);
+            match #fnname(py, &module) {
+                #krate::PyResult::Ok(()) => 0,
+                #krate::PyResult::Err(e) => {
+                    e.restore(py);
+                    -1
+                }
+            }
+        }
+
+        // An array of slot definitions for multi-phase initialization, terminated by a {0, NULL} entry.
+        #[doc(hidden)]
+        #visibility static #slots_name: [#krate::impl_::pymodule::ModuleDefSlot; 2] = [
+            #krate::impl_::pymodule::ModuleDefSlot::new(
+                #krate::ffi::Py_mod_exec,
+                #mod_exec_name as _,
+            ),
+            #krate::impl_::pymodule::ModuleDefSlot::new(
+                0,
+                ::std::ptr::null_mut(),
+            ),
+        ];
+
+        #[doc(hidden)]
+        #visibility static #module_def_name: #krate::impl_::pymodule::ModuleDef = unsafe {
+            #krate::impl_::pymodule::ModuleDef::new(
+                concat!(stringify!(#name), "\0"),
+                #doc,
+                &#slots_name as *const #krate::impl_::pymodule::ModuleDefSlot as *mut #krate::ffi::PyModuleDef_Slot,
+            )
+        };
+    };
+
+    if is_pypy() {
+        single_phase_module
+    } else if cfg!(feature = "pep489") {
+        multi_phase_module
+    } else {
+        single_phase_module
     }
 }
 
+fn is_pypy() -> bool {
+    let config = pyo3_build_config::get();
+    config.implementation.is_pypy()
+}
+
 /// 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();

pyo3-macros/Cargo.toml

@@ -18,6 +18,7 @@ multiple-pymethods = []
 
 pyproto = ["pyo3-macros-backend/pyproto"]
 abi3 = ["pyo3-macros-backend/abi3"]
+pep489 = ["pyo3-macros-backend/pep489"]
 
 [dependencies]
 proc-macro2 = { version = "1", default-features = false }

pytests/Cargo.toml

@@ -15,6 +15,9 @@ pyo3-build-config = { path = "../pyo3-build-config" }
 name = "pyo3_pytests"
 crate-type = ["cdylib"]
 
+[features]
+pep489 = ["pyo3/pep489"]
+
 [package.metadata.maturin]
 classifier=[
     "License :: OSI Approved :: MIT License",

pytests/noxfile.py

@@ -4,10 +4,13 @@
 
 
 @nox.session
-def test(session):
+@nox.parametrize("cargo_features", ['""', "pep489"])
+def test(session, cargo_features):
     session.install("-rrequirements-dev.txt")
     session.install("maturin")
-    session.run_always("maturin", "develop")
+    session.run_always(
+        "maturin", "develop", f"--cargo-extra-args=--features {cargo_features}"
+    )
     session.run("pytest")
 
 

src/impl_/pymodule.rs

@@ -1,23 +1,43 @@
 //! Implementation details of `#[pymodule]` which need to be accessible from proc-macro generated code.
 
 use std::cell::UnsafeCell;
+#[cfg(all(feature = "pep489", not(PyPy)))]
+use {std::ffi::CStr, std::os::raw::c_void};
 
-use crate::{
-    callback::panic_result_into_callback_output, ffi, types::PyModule, GILPool, IntoPyPointer, Py,
-    PyObject, PyResult, Python,
-};
+#[cfg(any(not(feature = "pep489"), PyPy))]
+use crate::{callback::panic_result_into_callback_output, GILPool, IntoPyPointer};
+use crate::{ffi, types::PyModule, Py, PyObject, PyResult, Python};
+#[cfg(all(feature = "pep489", not(PyPy)))]
+use crate::{AsPyPointer, PyErr};
+
+/// `Sync` wrapper of `ffi::PyModuleDef_Slot`
+#[cfg(all(feature = "pep489", not(PyPy)))]
+#[repr(transparent)]
+pub struct ModuleDefSlot(ffi::PyModuleDef_Slot);
+
+#[cfg(all(feature = "pep489", not(PyPy)))]
+unsafe impl Sync for ModuleDefSlot {}
 
 /// `Sync` wrapper of `ffi::PyModuleDef`.
 pub struct ModuleDef {
     // wrapped in UnsafeCell so that Rust compiler treats this as interior mutability
     ffi_def: UnsafeCell<ffi::PyModuleDef>,
+    #[cfg(any(not(feature = "pep489"), PyPy))]
     initializer: ModuleInitializer,
 }
 
+unsafe impl Sync for ModuleDef {}
+
 /// Wrapper to enable initializer to be used in const fns.
 pub struct ModuleInitializer(pub for<'py> fn(Python<'py>, &PyModule) -> PyResult<()>);
 
-unsafe impl Sync for ModuleDef {}
+#[cfg(all(feature = "pep489", not(PyPy)))]
+impl ModuleDefSlot {
+    /// Make new module definition slot
+    pub const fn new(slot: i32, value: *mut c_void) -> Self {
+        Self(ffi::PyModuleDef_Slot { slot, value })
+    }
+}
 
 impl ModuleDef {
     /// Make new module definition with given module name.
@@ -27,7 +47,8 @@ impl ModuleDef {
     pub const unsafe fn new(
         name: &'static str,
         doc: &'static str,
-        initializer: ModuleInitializer,
+        #[cfg(all(feature = "pep489", not(PyPy)))] slots: *mut ffi::PyModuleDef_Slot,
+        #[cfg(any(not(feature = "pep489"), PyPy))] initializer: ModuleInitializer,
     ) -> Self {
         const INIT: ffi::PyModuleDef = ffi::PyModuleDef {
             m_base: ffi::PyModuleDef_HEAD_INIT,
@@ -44,26 +65,62 @@ impl ModuleDef {
         let ffi_def = UnsafeCell::new(ffi::PyModuleDef {
             m_name: name.as_ptr() as *const _,
             m_doc: doc.as_ptr() as *const _,
+            #[cfg(all(feature = "pep489", not(PyPy)))]
+            m_slots: slots,
             ..INIT
         });
 
         ModuleDef {
             ffi_def,
+            #[cfg(any(not(feature = "pep489"), PyPy))]
             initializer,
         }
     }
+
+    /// Return module def
+    #[cfg(all(feature = "pep489", not(PyPy)))]
+    pub fn module_def(&'static self) -> *mut ffi::PyModuleDef {
+        self.ffi_def.get()
+    }
+
     /// Builds a module using user given initializer. Used for [`#[pymodule]`][crate::pymodule].
+    #[cfg(all(feature = "pep489", not(PyPy)))]
+    pub fn make_module(&'static self, py: Python<'_>) -> PyResult<PyObject> {
+        let module_spec_type = py.import("importlib.machinery")?.getattr("ModuleSpec")?;
+        let module = unsafe {
+            let mod_def = self.ffi_def.get();
+            // Mock a ModuleSpec object just good enough for PyModule_FromDefAndSpec()
+            // an object with just a name attribute.
+            let module_name = CStr::from_ptr((*mod_def).m_name);
+            let args = (module_name.to_str()?, py.None());
+            let spec = module_spec_type.call1(args)?;
+            let module = Py::<PyModule>::from_owned_ptr_or_err(
+                py,
+                ffi::PyModule_FromDefAndSpec(mod_def, spec.as_ptr()),
+            )?;
+            if ffi::PyModule_ExecDef(module.as_ptr(), mod_def) != 0 {
+                return Err(PyErr::fetch(py));
+            }
+            module
+        };
+        Ok(module.into())
+    }
+
+    /// Builds a module using user given initializer. Used for [`#[pymodule]`][crate::pymodule].
+    #[cfg(any(not(feature = "pep489"), PyPy))]
     pub fn make_module(&'static self, py: Python<'_>) -> PyResult<PyObject> {
         let module = unsafe {
             Py::<PyModule>::from_owned_ptr_or_err(py, ffi::PyModule_Create(self.ffi_def.get()))?
         };
         (self.initializer.0)(py, module.as_ref(py))?;
         Ok(module.into())
     }
+
     /// Implementation of `PyInit_foo` functions generated in [`#[pymodule]`][crate::pymodule]..
     ///
     /// # Safety
     /// The Python GIL must be held.
+    #[cfg(any(not(feature = "pep489"), PyPy))]
     pub unsafe fn module_init(&'static self) -> *mut ffi::PyObject {
         let pool = GILPool::new();
         let py = pool.python();
@@ -92,7 +149,7 @@ impl ModuleDef {
     }
 }
 
-#[cfg(test)]
+#[cfg(all(test, any(not(feature = "pep489"), PyPy)))]
 mod tests {
     use std::sync::atomic::{AtomicBool, Ordering};