support &Bound<PyType> for classmethod

pyo3-macros-backend/src/method.rs

@@ -127,13 +127,21 @@ impl FnType {
                 let slf: Ident = syn::Ident::new("_slf", Span::call_site());
                 quote_spanned! { *span =>
                     #[allow(clippy::useless_conversion)]
-                    ::std::convert::Into::into(_pyo3::types::PyType::from_type_ptr(#py, #slf.cast())),
+                    ::std::convert::Into::into(_pyo3::impl_::pymethods::BoundType(
+                        _pyo3::impl_::pymethods::ptr_to_bound(#py, &#slf.cast())
+                            .downcast_unchecked::<_pyo3::types::PyType>()
+                    )),
                 }
             }
             FnType::FnModule(span) => {
+                let py = syn::Ident::new("py", Span::call_site());
+                let slf: Ident = syn::Ident::new("_slf", Span::call_site());
                 quote_spanned! { *span =>
                     #[allow(clippy::useless_conversion)]
-                    ::std::convert::Into::into(py.from_borrowed_ptr::<_pyo3::types::PyModule>(_slf)),
+                    ::std::convert::Into::into(_pyo3::impl_::pymethods::BoundModule(
+                        _pyo3::impl_::pymethods::ptr_to_bound(#py, &#slf)
+                            .downcast_unchecked::<_pyo3::types::PyModule>()
+                    )),
                 }
             }
         }

pyo3-macros-backend/src/module.rs

@@ -93,7 +93,15 @@ pub fn pymodule_impl(
             use #krate::impl_::pymodule as impl_;
             impl #fnname::MakeDef {
                 const fn make_def() -> impl_::ModuleDef {
-                    const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(#fnname);
+                    /// Allow #[pymodule] to take any of &Bound<PyModule>, &PyModule, Py<PyModule>
+                    fn init_wrapper<'py>(py: Python<'py>, m: &Bound<'py, PyModule>) -> PyResult<()> {
+                        #fnname(
+                            py,
+                            ::std::convert::Into::into(::pyo3::impl_::pymethods::BoundModule(m))
+                        )
+                    }
+
+                    const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(init_wrapper);
                     unsafe {
                         impl_::ModuleDef::new(#fnname::NAME, #doc, INITIALIZER)
                     }

pytests/src/pyclasses.rs

@@ -44,6 +44,25 @@ struct AssertingBaseClass;
 
 #[pymethods]
 impl AssertingBaseClass {
+    #[new]
+    #[classmethod]
+    fn new(cls: &Bound<'_, PyType>, expected_type: Bound<'_, PyType>) -> PyResult<Self> {
+        if !cls.is(&expected_type) {
+            return Err(PyValueError::new_err(format!(
+                "{:?} != {:?}",
+                cls, expected_type
+            )));
+        }
+        Ok(Self)
+    }
+}
+
+#[pyclass(subclass)]
+#[derive(Clone, Debug)]
+struct AssertingBaseClassDeprecated;
+
+#[pymethods]
+impl AssertingBaseClassDeprecated {
     #[new]
     #[classmethod]
     fn new(cls: &PyType, expected_type: &PyType) -> PyResult<Self> {

src/derive_utils.rs

@@ -1,6 +1,6 @@
 //! Functionality for the code generated by the derive backend
 
-use crate::{types::PyModule, Python};
+use crate::{types::PyModule, Bound, Python};
 
 /// Enum to abstract over the arguments of Python function wrappers.
 pub enum PyFunctionArguments<'a> {
@@ -31,3 +31,33 @@ impl<'a> From<&'a PyModule> for PyFunctionArguments<'a> {
         PyFunctionArguments::PyModule(module)
     }
 }
+
+/// Enum to abstract over the arguments of Python function wrappers.
+pub enum PyFunctionArgumentsBound<'a, 'py> {
+    Python(Python<'py>),
+    PyModule(&'a Bound<'py, PyModule>),
+}
+
+impl<'a, 'py> PyFunctionArgumentsBound<'a, 'py> {
+    pub fn into_py_and_maybe_module(self) -> (Python<'py>, Option<&'a Bound<'py, PyModule>>) {
+        match self {
+            PyFunctionArgumentsBound::Python(py) => (py, None),
+            PyFunctionArgumentsBound::PyModule(module) => {
+                let py = module.py();
+                (py, Some(module))
+            }
+        }
+    }
+}
+
+impl<'a, 'py> From<Python<'py>> for PyFunctionArgumentsBound<'a, 'py> {
+    fn from(py: Python<'py>) -> PyFunctionArgumentsBound<'a, 'py> {
+        PyFunctionArgumentsBound::Python(py)
+    }
+}
+
+impl<'a, 'py> From<&'a Bound<'py, PyModule>> for PyFunctionArgumentsBound<'a, 'py> {
+    fn from(module: &'a Bound<'py, PyModule>) -> PyFunctionArgumentsBound<'a, 'py> {
+        PyFunctionArgumentsBound::PyModule(module)
+    }
+}

src/impl_/pyfunction.rs

@@ -1,4 +1,8 @@
-use crate::{derive_utils::PyFunctionArguments, types::PyCFunction, PyResult};
+use crate::{
+    derive_utils::{PyFunctionArguments, PyFunctionArgumentsBound},
+    types::PyCFunction,
+    Bound, PyResult,
+};
 
 pub use crate::impl_::pymethods::PyMethodDef;
 
@@ -8,3 +12,10 @@ pub fn _wrap_pyfunction<'a>(
 ) -> PyResult<&'a PyCFunction> {
     PyCFunction::internal_new(method_def, py_or_module.into())
 }
+
+pub fn _wrap_pyfunction_bound<'a, 'py: 'a>(
+    method_def: &PyMethodDef,
+    py_or_module: impl Into<PyFunctionArgumentsBound<'a, 'py>>,
+) -> PyResult<Bound<'py, PyCFunction>> {
+    PyCFunction::internal_new_bound(method_def, py_or_module.into())
+}

src/impl_/pymethods.rs

@@ -3,8 +3,10 @@ use crate::exceptions::PyStopAsyncIteration;
 use crate::gil::LockGIL;
 use crate::impl_::panic::PanicTrap;
 use crate::internal_tricks::extract_c_string;
+use crate::types::{PyModule, PyType};
 use crate::{
-    ffi, PyAny, PyCell, PyClass, PyErr, PyObject, PyResult, PyTraverseError, PyVisit, Python,
+    ffi, Bound, Py, PyAny, PyCell, PyClass, PyErr, PyObject, PyResult, PyTraverseError, PyVisit,
+    Python,
 };
 use std::borrow::Cow;
 use std::ffi::CStr;
@@ -466,3 +468,62 @@ pub trait AsyncIterResultOptionKind {
 }
 
 impl<Value, Error> AsyncIterResultOptionKind for Result<Option<Value>, Error> {}
+
+/// Create a reference to bound from a reference to a raw ffi pointer
+///
+/// # Safety:
+///   - ptr must be a non-null pointer to a valid ffi::PyObject
+///   - ptr owns a Python reference for at least the lifetime 'a
+#[inline]
+pub unsafe fn ptr_to_bound<'py, 'a>(
+    py: Python<'py>,
+    ptr: &'a *mut ffi::PyObject,
+) -> &'a Bound<'py, PyAny> {
+    Bound::from_ref_to_ptr(py, ptr)
+}
+
+pub struct BoundType<'a, 'py>(pub &'a Bound<'py, PyType>);
+
+impl<'a> From<BoundType<'a, 'a>> for &'a PyType {
+    #[inline]
+    fn from(bound: BoundType<'a, 'a>) -> Self {
+        bound.0.as_gil_ref()
+    }
+}
+
+impl<'a, 'py> From<BoundType<'a, 'py>> for &'a Bound<'py, PyType> {
+    #[inline]
+    fn from(bound: BoundType<'a, 'py>) -> Self {
+        bound.0
+    }
+}
+
+impl From<BoundType<'_, '_>> for Py<PyType> {
+    #[inline]
+    fn from(bound: BoundType<'_, '_>) -> Self {
+        bound.0.clone().unbind()
+    }
+}
+
+pub struct BoundModule<'a, 'py>(pub &'a Bound<'py, PyModule>);
+
+impl<'a> From<BoundModule<'a, 'a>> for &'a PyModule {
+    #[inline]
+    fn from(bound: BoundModule<'a, 'a>) -> Self {
+        bound.0.as_gil_ref()
+    }
+}
+
+impl<'a, 'py> From<BoundModule<'a, 'py>> for &'a Bound<'py, PyModule> {
+    #[inline]
+    fn from(bound: BoundModule<'a, 'py>) -> Self {
+        bound.0
+    }
+}
+
+impl From<BoundModule<'_, '_>> for Py<PyModule> {
+    #[inline]
+    fn from(bound: BoundModule<'_, '_>) -> Self {
+        bound.0.clone().unbind()
+    }
+}

src/impl_/pymodule.rs

@@ -7,7 +7,7 @@ use portable_atomic::{AtomicI64, Ordering};
 
 #[cfg(not(PyPy))]
 use crate::exceptions::PyImportError;
-use crate::{ffi, sync::GILOnceCell, types::PyModule, Py, PyResult, Python};
+use crate::{ffi, sync::GILOnceCell, types::PyModule, Bound, Py, PyResult, Python};
 
 /// `Sync` wrapper of `ffi::PyModuleDef`.
 pub struct ModuleDef {
@@ -22,7 +22,7 @@ pub struct ModuleDef {
 }
 
 /// Wrapper to enable initializer to be used in const fns.
-pub struct ModuleInitializer(pub for<'py> fn(Python<'py>, &PyModule) -> PyResult<()>);
+pub struct ModuleInitializer(pub for<'py> fn(Python<'py>, &Bound<'py, PyModule>) -> PyResult<()>);
 
 unsafe impl Sync for ModuleDef {}
 
@@ -125,7 +125,7 @@ impl ModuleDef {
                         ffi::PyModule_Create(self.ffi_def.get()),
                     )?
                 };
-                (self.initializer.0)(py, module.as_ref(py))?;
+                (self.initializer.0)(py, module.bind(py))?;
                 Ok(module)
             })
             .map(|py_module| py_module.clone_ref(py))
@@ -136,7 +136,9 @@ impl ModuleDef {
 mod tests {
     use std::sync::atomic::{AtomicBool, Ordering};
 
-    use crate::{types::PyModule, PyResult, Python};
+    use crate::{
+        types::module::PyModuleMethods, types::PyModule, Bound, PyNativeType, PyResult, Python,
+    };
 
     use super::{ModuleDef, ModuleInitializer};
 
@@ -191,7 +193,7 @@ mod tests {
         static INIT_CALLED: AtomicBool = AtomicBool::new(false);
 
         #[allow(clippy::unnecessary_wraps)]
-        fn init(_: Python<'_>, _: &PyModule) -> PyResult<()> {
+        fn init(_: Python<'_>, _: &Bound<'_, PyModule>) -> PyResult<()> {
             INIT_CALLED.store(true, Ordering::SeqCst);
             Ok(())
         }
@@ -202,7 +204,8 @@ mod tests {
             assert_eq!((*module_def.ffi_def.get()).m_doc, DOC.as_ptr() as _);
 
             Python::with_gil(|py| {
-                module_def.initializer.0(py, py.import("builtins").unwrap()).unwrap();
+                module_def.initializer.0(py, &py.import("builtins").unwrap().as_borrowed())
+                    .unwrap();
                 assert!(INIT_CALLED.load(Ordering::SeqCst));
             })
         }

src/instance.rs

@@ -94,6 +94,14 @@ impl<'py> Bound<'py, PyAny> {
     ) -> PyResult<Self> {
         Py::from_owned_ptr_or_err(py, ptr).map(|obj| Self(py, ManuallyDrop::new(obj)))
     }
+
+    #[inline]
+    pub(crate) unsafe fn from_ref_to_ptr<'a>(
+        _py: Python<'py>,
+        ptr: &'a *mut ffi::PyObject,
+    ) -> &'a Self {
+        &*(ptr as *const *mut ffi::PyObject).cast::<Bound<'py, PyAny>>()
+    }
 }
 
 impl<'py, T> Bound<'py, T> {

src/macros.rs

@@ -134,6 +134,27 @@ macro_rules! wrap_pyfunction {
     }};
 }
 
+/// Wraps a Rust function annotated with [`#[pyfunction]`](macro@crate::pyfunction).
+///
+/// This can be used with [`PyModule::add_function`](crate::types::PyModule::add_function) to add free
+/// functions to a [`PyModule`](crate::types::PyModule) - see its documentation for more information.
+#[macro_export]
+macro_rules! wrap_pyfunction_bound {
+    ($function:path) => {
+        &|py_or_module| {
+            use $function as wrapped_pyfunction;
+            $crate::impl_::pyfunction::_wrap_pyfunction_bound(
+                &wrapped_pyfunction::DEF,
+                py_or_module,
+            )
+        }
+    };
+    ($function:path, $py_or_module:expr) => {{
+        use $function as wrapped_pyfunction;
+        $crate::impl_::pyfunction::_wrap_pyfunction_bound(&wrapped_pyfunction::DEF, $py_or_module)
+    }};
+}
+
 /// Returns a function that takes a [`Python`](crate::Python) instance and returns a
 /// Python module.
 ///

src/prelude.rs

@@ -23,7 +23,7 @@ pub use crate::PyNativeType;
 pub use pyo3_macros::{pyclass, pyfunction, pymethods, pymodule, FromPyObject};
 
 #[cfg(feature = "macros")]
-pub use crate::wrap_pyfunction;
+pub use crate::{wrap_pyfunction, wrap_pyfunction_bound};
 
 pub use crate::types::any::PyAnyMethods;
 pub use crate::types::boolobject::PyBoolMethods;

src/types/function.rs

@@ -1,6 +1,8 @@
-use crate::derive_utils::PyFunctionArguments;
+use crate::derive_utils::{PyFunctionArguments, PyFunctionArgumentsBound};
+use crate::ffi_ptr_ext::FfiPtrExt;
 use crate::methods::PyMethodDefDestructor;
 use crate::prelude::*;
+use crate::py_result_ext::PyResultExt;
 use crate::{
     ffi,
     impl_::pymethods::{self, PyMethodDef},
@@ -132,6 +134,35 @@ impl PyCFunction {
             ))
         }
     }
+
+    #[doc(hidden)]
+    pub(crate) fn internal_new_bound<'a, 'py>(
+        method_def: &PyMethodDef,
+        py_or_module: PyFunctionArgumentsBound<'a, 'py>,
+    ) -> PyResult<Bound<'py, Self>> {
+        let (py, module) = py_or_module.into_py_and_maybe_module();
+        let (mod_ptr, module_name): (_, Option<Py<PyString>>) = if let Some(m) = module {
+            let mod_ptr = m.as_ptr();
+            (mod_ptr, Some(m.name()?.into_py(py)))
+        } else {
+            (std::ptr::null_mut(), None)
+        };
+        let (def, destructor) = method_def.as_method_def()?;
+
+        // FIXME: stop leaking the def and destructor
+        let def = Box::into_raw(Box::new(def));
+        std::mem::forget(destructor);
+
+        let module_name_ptr = module_name
+            .as_ref()
+            .map_or(std::ptr::null_mut(), Py::as_ptr);
+
+        unsafe {
+            ffi::PyCFunction_NewEx(def, mod_ptr, module_name_ptr)
+                .assume_owned_or_err(py)
+                .downcast_into_unchecked()
+        }
+    }
 }
 
 fn closure_capsule_name() -> &'static CStr {