From 1bd13386c66fc79738a6e6e54cc60f02d1cd6592 Mon Sep 17 00:00:00 2001
From: Nathan Goldbaum <nathan.goldbaum@gmail.com>
Date: Mon, 21 Oct 2024 12:48:27 -0600
Subject: [PATCH] add a way to declare free-threaded support without macros

---
 src/impl_/pymodule.rs | 20 ++++++++++++-------
 src/types/module.rs   | 45 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/src/impl_/pymodule.rs b/src/impl_/pymodule.rs
index 92949a30108..cdbe18d5e5e 100644
--- a/src/impl_/pymodule.rs
+++ b/src/impl_/pymodule.rs
@@ -1,6 +1,6 @@
 //! Implementation details of `#[pymodule]` which need to be accessible from proc-macro generated code.
 
-use std::{cell::UnsafeCell, ffi::CStr, marker::PhantomData};
+use std::{cell::UnsafeCell, ffi::CStr, marker::PhantomData, os::raw::c_int};
 
 #[cfg(all(
     not(any(PyPy, GraalPy)),
@@ -24,7 +24,7 @@ use crate::{
     impl_::pymethods::PyMethodDef,
     sync::GILOnceCell,
     types::{PyCFunction, PyModule, PyModuleMethods},
-    Bound, Py, PyClass, PyResult, PyTypeInfo, Python,
+    Bound, Py, PyClass, PyErr, PyResult, PyTypeInfo, Python,
 };
 
 /// `Sync` wrapper of `ffi::PyModuleDef`.
@@ -141,11 +141,17 @@ impl ModuleDef {
                         ffi::PyModule_Create(self.ffi_def.get()),
                     )?
                 };
-                if supports_free_threaded {
-                    unsafe {
-                        ffi::PyUnstable_Module_SetGIL(module.as_ptr(), ffi::Py_MOD_GIL_NOT_USED)
-                    };
-                }
+                let gil_used = {
+                    if supports_free_threaded {
+                        ffi::Py_MOD_GIL_NOT_USED
+                    } else {
+                        ffi::Py_MOD_GIL_USED
+                    }
+                };
+                match unsafe { ffi::PyUnstable_Module_SetGIL(module.as_ptr(), gil_used) } {
+                    c_int::MIN..=-1 => return Err(PyErr::fetch(py)),
+                    0..=c_int::MAX => {}
+                };
                 self.initializer.0(module.bind(py))?;
                 Ok(module)
             })
diff --git a/src/types/module.rs b/src/types/module.rs
index c0ed2df70c7..0854eb483bc 100644
--- a/src/types/module.rs
+++ b/src/types/module.rs
@@ -9,6 +9,7 @@ use crate::types::{
 };
 use crate::{exceptions, ffi, Borrowed, Bound, BoundObject, Py, PyObject, Python};
 use std::ffi::{CStr, CString};
+use std::os::raw::c_int;
 use std::str;
 
 /// Represents a Python [`module`][1] object.
@@ -384,6 +385,38 @@ pub trait PyModuleMethods<'py>: crate::sealed::Sealed {
     /// [1]: crate::prelude::pyfunction
     /// [2]: crate::wrap_pyfunction
     fn add_function(&self, fun: Bound<'_, PyCFunction>) -> PyResult<()>;
+
+    /// Declare whether or not this module supports running with the GIL disabled
+    ///
+    /// If the module does not rely on the GIL for thread safety, you can pass True
+    /// to this function so that when the module is imported the interpreter will
+    /// not enable the GIL at runtime on the free-threaded interpreter.
+    ///
+    /// This function sets the [`Py_MOD_GIL`
+    /// slot](https://docs.python.org/3/c-api/module.html#c.Py_mod_gil) on the
+    /// module object. The default is `Py_MOD_GIL_USED`, so passing `false` to
+    /// this function is a no-op unless you have already set `Py_MOD_GIL` to
+    /// `Py_MOD_GIL_NOT_USED` elsewhere.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use pyo3::prelude::*;
+    ///
+    /// #[pymodule(supports_free_threaded = true)]
+    /// fn my_module(py: Python<'_>, module: &Bound<'_, PyModule>) -> PyResult<()> {
+    ///     let submodule = PyModule::new(py, "submodule")?;
+    ///     submodule.supports_free_threaded(true)?;
+    ///     module.add_submodule(&submodule)?;
+    ///     Ok(())
+    /// }
+    /// ```
+    ///
+    /// The resulting module will not print a `RuntimeWarning` and re-enable the
+    /// GIL when Python imports it on the free-threaded build, since all module
+    /// objects defined in the extension have `Py_MOD_GIL` set to
+    /// `Py_MOD_GIL_NOT_USED`.
+    fn supports_free_threaded(&self, supports_free_threaded: bool) -> PyResult<()>;
 }
 
 impl<'py> PyModuleMethods<'py> for Bound<'py, PyModule> {
@@ -493,7 +526,6 @@ impl<'py> PyModuleMethods<'py> for Bound<'py, PyModule> {
         T: IntoPyCallbackOutput<'py, PyObject>,
     {
         fn inner(module: &Bound<'_, PyModule>, object: Bound<'_, PyAny>) -> PyResult<()> {
-            if object.is_instance_of::<PyModule>() {}
             let name = object.getattr(__name__(module.py()))?;
             module.add(name.downcast_into::<PyString>()?, object)
         }
@@ -511,6 +543,17 @@ impl<'py> PyModuleMethods<'py> for Bound<'py, PyModule> {
         let name = fun.getattr(__name__(self.py()))?;
         self.add(name.downcast_into::<PyString>()?, fun)
     }
+
+    fn supports_free_threaded(&self, supports_free_threaded: bool) -> PyResult<()> {
+        let gil_used = match supports_free_threaded {
+            true => ffi::Py_MOD_GIL_NOT_USED,
+            false => ffi::Py_MOD_GIL_USED,
+        };
+        match unsafe { ffi::PyUnstable_Module_SetGIL(self.as_ptr(), gil_used) } {
+            c_int::MIN..=-1 => Err(PyErr::fetch(self.py())),
+            0..=c_int::MAX => Ok(()),
+        }
+    }
 }
 
 fn __all__(py: Python<'_>) -> &Bound<'_, PyString> {