relax multiple-import check to only prevent subinterpreters

Author: davidhewitt

newsfragments/3447.changed.md

@@ -0,0 +1 @@
+Relax multiple import check to only prevent use of subinterpreters.

pyo3-build-config/src/lib.rs

@@ -146,6 +146,11 @@ pub fn print_feature_cfgs() {
     if rustc_minor_version >= 59 {
         println!("cargo:rustc-cfg=thread_local_const_init");
     }
+
+    // Enable use of OnceLock on Rust 1.70 and greater
+    if rustc_minor_version >= 70 {
+        println!("cargo:rustc-cfg=once_lock");
+    }
 }
 
 /// Private exports used in PyO3's build.rs

pytests/tests/test_misc.py

@@ -10,15 +10,27 @@ def test_issue_219():
     pyo3_pytests.misc.issue_219()
 
 
+def test_multiple_imports_same_interpreter_ok():
+    spec = importlib.util.find_spec("pyo3_pytests.pyo3_pytests")
+
+    module = importlib.util.module_from_spec(spec)
+    assert dir(module) == dir(pyo3_pytests.pyo3_pytests)
+
+
 @pytest.mark.skipif(
     platform.python_implementation() == "PyPy",
-    reason="PyPy does not reinitialize the module (appears to be some internal caching)",
+    reason="PyPy does not support subinterpreters",
 )
-def test_second_module_import_fails():
-    spec = importlib.util.find_spec("pyo3_pytests.pyo3_pytests")
+def test_import_in_subinterpreter_forbidden():
+    import _xxsubinterpreters
 
+    sub_interpreter = _xxsubinterpreters.create()
     with pytest.raises(
-        ImportError,
-        match="PyO3 modules may only be initialized once per interpreter process",
+        _xxsubinterpreters.RunFailedError,
+        match="PyO3 modules do not yet support subinterpreters, see https://github.com/PyO3/pyo3/issues/576",
     ):
-        importlib.util.module_from_spec(spec)
+        _xxsubinterpreters.run_string(
+            sub_interpreter, "import pyo3_pytests.pyo3_pytests"
+        )
+
+    _xxsubinterpreters.destroy(sub_interpreter)

src/impl_/pymodule.rs

@@ -1,9 +1,11 @@
 //! Implementation details of `#[pymodule]` which need to be accessible from proc-macro generated code.
 
-use std::{
-    cell::UnsafeCell,
-    sync::atomic::{self, AtomicBool},
-};
+use std::cell::UnsafeCell;
+#[cfg(once_lock)]
+use std::sync::OnceLock;
+
+#[cfg(not(once_lock))]
+use parking_lot::Mutex;
 
 use crate::{exceptions::PyImportError, ffi, types::PyModule, Py, PyResult, Python};
 
@@ -12,7 +14,10 @@ pub struct ModuleDef {
     // wrapped in UnsafeCell so that Rust compiler treats this as interior mutability
     ffi_def: UnsafeCell<ffi::PyModuleDef>,
     initializer: ModuleInitializer,
-    initialized: AtomicBool,
+    #[cfg(once_lock)]
+    interpreter: OnceLock<i64>,
+    #[cfg(not(once_lock))]
+    interpreter: Mutex<Option<i64>>,
 }
 
 /// Wrapper to enable initializer to be used in const fns.
@@ -51,7 +56,10 @@ impl ModuleDef {
         ModuleDef {
             ffi_def,
             initializer,
-            initialized: AtomicBool::new(false),
+            #[cfg(once_lock)]
+            interpreter: OnceLock::new(),
+            #[cfg(not(once_lock))]
+            interpreter: Mutex::new(None),
         }
     }
     /// Builds a module using user given initializer. Used for [`#[pymodule]`][crate::pymodule].
@@ -74,9 +82,22 @@ impl ModuleDef {
         let module = unsafe {
             Py::<PyModule>::from_owned_ptr_or_err(py, ffi::PyModule_Create(self.ffi_def.get()))?
         };
-        if self.initialized.swap(true, atomic::Ordering::SeqCst) {
+        let current_interpreter =
+            unsafe { ffi::PyInterpreterState_GetID(ffi::PyInterpreterState_Get()) };
+        let initialized_interpreter = py.allow_threads(|| {
+            #[cfg(once_lock)]
+            {
+                *self.interpreter.get_or_init(|| current_interpreter)
+            }
+
+            #[cfg(not(once_lock))]
+            {
+                *self.interpreter.lock().get_or_insert(current_interpreter)
+            }
+        });
+        if current_interpreter != initialized_interpreter {
             return Err(PyImportError::new_err(
-                "PyO3 modules may only be initialized once per interpreter process",
+                "PyO3 modules do not yet support subinterpreters, see https://github.com/PyO3/pyo3/issues/576",
             ));
         }
         (self.initializer.0)(py, module.as_ref(py))?;