From 0e430716af8e77ff97111629502854812a00ff05 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bla=C5=BE=20=C5=A0nuderl?= <snuderl@gmail.com>
Date: Mon, 5 Feb 2024 14:21:10 +0100
Subject: [PATCH] Make #wrap_pyfunction return a Bound value, and update
 module.add_method to accept it

---
 src/impl_/pyfunction.rs                 | 8 ++++----
 src/sync.rs                             | 1 +
 src/types/module.rs                     | 2 +-
 tests/test_coroutine.rs                 | 5 ++++-
 tests/test_wrap_pyfunction_deduction.rs | 2 +-
 5 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/impl_/pyfunction.rs b/src/impl_/pyfunction.rs
index 464beeb8ce5..c655b97b6e2 100644
--- a/src/impl_/pyfunction.rs
+++ b/src/impl_/pyfunction.rs
@@ -1,10 +1,10 @@
-use crate::{derive_utils::PyFunctionArguments, types::PyCFunction, PyResult};
-
 pub use crate::impl_::pymethods::PyMethodDef;
+use crate::Bound;
+use crate::{derive_utils::PyFunctionArguments, types::PyCFunction, PyResult};
 
 pub fn _wrap_pyfunction<'a>(
     method_def: &PyMethodDef,
     py_or_module: impl Into<PyFunctionArguments<'a>>,
-) -> PyResult<&'a PyCFunction> {
-    PyCFunction::internal_new(method_def, py_or_module.into()).map(|x| x.into_gil_ref())
+) -> PyResult<Bound<'a, PyCFunction>> {
+    PyCFunction::internal_new(method_def, py_or_module.into())
 }
diff --git a/src/sync.rs b/src/sync.rs
index a88a0f4b485..ea93fd75e75 100644
--- a/src/sync.rs
+++ b/src/sync.rs
@@ -211,6 +211,7 @@ impl GILOnceCell<Py<PyType>> {
 /// ```
 /// use pyo3::intern;
 /// # use pyo3::{pyfunction, types::PyDict, wrap_pyfunction, PyResult, Python};
+/// use crate::pyo3::prelude::PyAnyMethods;
 ///
 /// #[pyfunction]
 /// fn create_dict(py: Python<'_>) -> PyResult<&PyDict> {
diff --git a/src/types/module.rs b/src/types/module.rs
index 8824dfcf030..38d8256ac76 100644
--- a/src/types/module.rs
+++ b/src/types/module.rs
@@ -348,7 +348,7 @@ impl PyModule {
     ///
     /// [1]: crate::prelude::pyfunction
     /// [2]: crate::wrap_pyfunction
-    pub fn add_function<'a>(&'a self, fun: &'a PyCFunction) -> PyResult<()> {
+    pub fn add_function<'a>(&'a self, fun: Bound<'a, PyCFunction>) -> PyResult<()> {
         self.as_borrowed().add_function(&fun.as_borrowed())
     }
 }
diff --git a/tests/test_coroutine.rs b/tests/test_coroutine.rs
index 206c35da93c..6ec04f144c1 100644
--- a/tests/test_coroutine.rs
+++ b/tests/test_coroutine.rs
@@ -65,7 +65,10 @@ fn test_coroutine_qualname() {
             assert coro.__name__ == name and coro.__qualname__ == qualname
         "#;
         let locals = [
-            ("my_fn", wrap_pyfunction!(my_fn, gil).unwrap().deref()),
+            (
+                "my_fn",
+                wrap_pyfunction!(my_fn, gil).unwrap().into_gil_ref().deref(),
+            ),
             ("MyClass", gil.get_type::<MyClass>()),
         ]
         .into_py_dict(gil);
diff --git a/tests/test_wrap_pyfunction_deduction.rs b/tests/test_wrap_pyfunction_deduction.rs
index 8723ad43fbd..638b980810c 100644
--- a/tests/test_wrap_pyfunction_deduction.rs
+++ b/tests/test_wrap_pyfunction_deduction.rs
@@ -5,7 +5,7 @@ use pyo3::{prelude::*, types::PyCFunction};
 #[pyfunction]
 fn f() {}
 
-pub fn add_wrapped(wrapper: &impl Fn(Python<'_>) -> PyResult<&PyCFunction>) {
+pub fn add_wrapped(wrapper: &impl Fn(Python<'_>) -> PyResult<Bound<'_, PyCFunction>>) {
     let _ = wrapper;
 }