Skip to content

Commit

Permalink
feat: allow classmethod/pass_module to receive owned types
Browse files Browse the repository at this point in the history
This is necessary for async functions
  • Loading branch information
Joseph Perez authored and wyfo committed Nov 22, 2023
1 parent c8fdb80 commit 744de3a
Show file tree
Hide file tree
Showing 6 changed files with 58 additions and 13 deletions.
2 changes: 2 additions & 0 deletions newsfragments/3587.added.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
- Classmethods can now receive `Py<PyType>` as their first argument
- Function annotated with `pass_module` can now receive `Py<PyModule>` as their first argument
6 changes: 4 additions & 2 deletions pyo3-macros-backend/src/method.rs
Original file line number Diff line number Diff line change
Expand Up @@ -113,12 +113,14 @@ impl FnType {
}
FnType::FnClass | FnType::FnNewClass => {
quote! {
_pyo3::types::PyType::from_type_ptr(py, _slf as *mut _pyo3::ffi::PyTypeObject),
#[allow(clippy::useless_conversion)]
::std::convert::Into::into(_pyo3::types::PyType::from_type_ptr(py, _slf as *mut _pyo3::ffi::PyTypeObject)),
}
}
FnType::FnModule => {
quote! {
py.from_borrowed_ptr::<_pyo3::types::PyModule>(_slf),
#[allow(clippy::useless_conversion)]
::std::convert::Into::into(py.from_borrowed_ptr::<_pyo3::types::PyModule>(_slf)),
}
}
}
Expand Down
35 changes: 25 additions & 10 deletions pyo3-macros-backend/src/pyfunction.rs
Original file line number Diff line number Diff line change
Expand Up @@ -199,7 +199,8 @@ pub fn impl_wrap_pyfunction(
.collect::<syn::Result<Vec<_>>>()?;

let tp = if pass_module.is_some() {
const PASS_MODULE_ERR: &str = "expected &PyModule as first argument with `pass_module`";
const PASS_MODULE_ERR: &str =
"expected &PyModule or Py<PyModule> as first argument with `pass_module`";
ensure_spanned!(
!arguments.is_empty(),
func.span() => PASS_MODULE_ERR
Expand Down Expand Up @@ -271,18 +272,32 @@ pub fn impl_wrap_pyfunction(
}

fn type_is_pymodule(ty: &syn::Type) -> bool {
if let syn::Type::Reference(tyref) = ty {
if let syn::Type::Path(typath) = tyref.elem.as_ref() {
if typath
.path
.segments
.last()
.map(|seg| seg.ident == "PyModule")
.unwrap_or(false)
let is_pymodule = |typath: &syn::TypePath| {
typath
.path
.segments
.last()
.map_or(false, |seg| seg.ident == "PyModule")
};
match ty {
syn::Type::Reference(tyref) => {
if let syn::Type::Path(typath) = tyref.elem.as_ref() {
return is_pymodule(typath);
}
}
syn::Type::Path(typath) => {
if let Some(syn::PathSegment {
arguments: syn::PathArguments::AngleBracketed(args),
..
}) = typath.path.segments.last()
{
return true;
if args.args.len() != 1 {
return false;
}
return matches!(args.args.first().unwrap(), syn::GenericArgument::Type(syn::Type::Path(typath)) if is_pymodule(typath));
}
}
_ => {}
}
false
}
13 changes: 13 additions & 0 deletions tests/test_methods.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,14 @@ impl ClassMethod {
fn method(cls: &PyType) -> PyResult<String> {
Ok(format!("{}.method()!", cls.name()?))
}

#[classmethod]
fn method_owned(cls: Py<PyType>) -> PyResult<String> {
Ok(format!(
"{}.method_owned()!",
Python::with_gil(|gil| cls.as_ref(gil).name().map(ToString::to_string))?
))
}
}

#[test]
Expand All @@ -84,6 +92,11 @@ fn class_method() {
let d = [("C", py.get_type::<ClassMethod>())].into_py_dict(py);
py_assert!(py, *d, "C.method() == 'ClassMethod.method()!'");
py_assert!(py, *d, "C().method() == 'ClassMethod.method()!'");
py_assert!(
py,
*d,
"C().method_owned() == 'ClassMethod.method_owned()!'"
);
py_assert!(py, *d, "C.method.__doc__ == 'Test class method.'");
py_assert!(py, *d, "C().method.__doc__ == 'Test class method.'");
});
Expand Down
13 changes: 13 additions & 0 deletions tests/test_module.rs
Original file line number Diff line number Diff line change
Expand Up @@ -348,6 +348,12 @@ fn pyfunction_with_module(module: &PyModule) -> PyResult<&str> {
module.name()
}

#[pyfunction]
#[pyo3(pass_module)]
fn pyfunction_with_module_owned(module: Py<PyModule>) -> PyResult<String> {
Python::with_gil(|gil| module.as_ref(gil).name().map(Into::into))
}

#[pyfunction]
#[pyo3(pass_module)]
fn pyfunction_with_module_and_py<'a>(
Expand Down Expand Up @@ -393,6 +399,7 @@ fn pyfunction_with_pass_module_in_attribute(module: &PyModule) -> PyResult<&str>
#[pymodule]
fn module_with_functions_with_module(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(pyfunction_with_module, m)?)?;
m.add_function(wrap_pyfunction!(pyfunction_with_module_owned, m)?)?;
m.add_function(wrap_pyfunction!(pyfunction_with_module_and_py, m)?)?;
m.add_function(wrap_pyfunction!(pyfunction_with_module_and_arg, m)?)?;
m.add_function(wrap_pyfunction!(pyfunction_with_module_and_default_arg, m)?)?;
Expand All @@ -401,6 +408,7 @@ fn module_with_functions_with_module(_py: Python<'_>, m: &PyModule) -> PyResult<
pyfunction_with_pass_module_in_attribute,
m
)?)?;
m.add_function(wrap_pyfunction!(pyfunction_with_module, m)?)?;
Ok(())
}

Expand All @@ -413,6 +421,11 @@ fn test_module_functions_with_module() {
m,
"m.pyfunction_with_module() == 'module_with_functions_with_module'"
);
py_assert!(
py,
m,
"m.pyfunction_with_module_owned() == 'module_with_functions_with_module'"
);
py_assert!(
py,
m,
Expand Down
2 changes: 1 addition & 1 deletion tests/ui/invalid_need_module_arg_position.stderr
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
error: expected &PyModule as first argument with `pass_module`
error: expected &PyModule or Py<PyModule> as first argument with `pass_module`
--> tests/ui/invalid_need_module_arg_position.rs:6:21
|
6 | fn fail(string: &str, module: &PyModule) -> PyResult<&str> {
Expand Down

0 comments on commit 744de3a

Please sign in to comment.