From 6b95d0d3c973eeb50174087201b6dba30f774127 Mon Sep 17 00:00:00 2001 From: Joseph Perez Date: Mon, 20 Nov 2023 10:39:05 +0100 Subject: [PATCH] feat: add coroutine `__name__`/`__qualname__` and not-awaited warning --- newsfragments/3588.added.md | 1 + pyo3-macros-backend/src/method.rs | 18 ++++++++- src/coroutine.rs | 47 ++++++++++++++++++++-- src/impl_/coroutine.rs | 41 +++++++++++++------ src/impl_/wrap.rs | 24 ++++++++++++ tests/test_coroutine.rs | 65 +++++++++++++++++++++++++++++++ 6 files changed, 180 insertions(+), 16 deletions(-) create mode 100644 newsfragments/3588.added.md diff --git a/newsfragments/3588.added.md b/newsfragments/3588.added.md new file mode 100644 index 00000000000..20c9a3c7bd4 --- /dev/null +++ b/newsfragments/3588.added.md @@ -0,0 +1 @@ +Add `__name__`/`__qualname__` attributes to `Coroutine`, as well as a Python warning when the coroutine is dropped without having been awaited \ No newline at end of file diff --git a/pyo3-macros-backend/src/method.rs b/pyo3-macros-backend/src/method.rs index 034d079bb46..a84f109f3f4 100644 --- a/pyo3-macros-backend/src/method.rs +++ b/pyo3-macros-backend/src/method.rs @@ -451,7 +451,23 @@ impl<'a> FnSpec<'a> { let rust_call = |args: Vec| { let mut call = quote! { function(#self_arg #(#args),*) }; if self.asyncness.is_some() { - call = quote! { _pyo3::impl_::coroutine::wrap_future(#call) }; + let python_name = &self.python_name; + let qualname = match cls { + Some(cls) => quote! { + _pyo3::impl_::coroutine::method_coroutine_qualname::<#cls>(py, stringify!(#python_name)) + }, + None => quote! { + _pyo3::impl_::coroutine::coroutine_qualname(py, py.from_borrowed_ptr_or_opt::<_pyo3::types::PyModule>(_slf), stringify!(#python_name)) + }, + }; + call = quote! {{ + let future = #call; + _pyo3::impl_::coroutine::new_coroutine( + _pyo3::types::PyString::new(py, stringify!(#python_name)).into(), + #qualname, + async move { _pyo3::impl_::wrap::OkWrap::wrap_no_gil(future.await) } + ) + }}; } quotes::map_result_into_ptr(quotes::ok_wrap(call)) }; diff --git a/src/coroutine.rs b/src/coroutine.rs index 564262f3bf4..5250aefd47c 100644 --- a/src/coroutine.rs +++ b/src/coroutine.rs @@ -14,10 +14,10 @@ use pyo3_macros::{pyclass, pymethods}; use crate::{ coroutine::waker::AsyncioWaker, - exceptions::{PyRuntimeError, PyStopIteration}, + exceptions::{PyAttributeError, PyRuntimeError, PyRuntimeWarning, PyStopIteration}, panic::PanicException, pyclass::IterNextOutput, - types::PyIterator, + types::{PyIterator, PyString}, IntoPy, Py, PyAny, PyErr, PyObject, PyResult, Python, }; @@ -30,6 +30,8 @@ type FutureOutput = Result, Box>; /// Python coroutine wrapping a [`Future`]. #[pyclass(crate = "crate")] pub struct Coroutine { + name: Option>, + qualname: Option>, future: Option + Send>>>, waker: Option>, } @@ -41,7 +43,11 @@ impl Coroutine { /// (should always be `None` anyway). /// /// `Coroutine `throw` drop the wrapped future and reraise the exception passed - pub(crate) fn from_future(future: F) -> Self + pub(crate) fn new( + name: Option>, + mut qualname: Option>, + future: F, + ) -> Self where F: Future> + Send + 'static, T: IntoPy, @@ -52,7 +58,10 @@ impl Coroutine { // SAFETY: GIL is acquired when future is polled (see `Coroutine::poll`) Ok(obj.into_py(unsafe { Python::assume_gil_acquired() })) }; + qualname = qualname.or_else(|| name.clone()); Self { + name, + qualname, future: Some(Box::pin(panic::AssertUnwindSafe(wrap).catch_unwind())), waker: None, } @@ -113,6 +122,20 @@ pub(crate) fn iter_result(result: IterNextOutput) -> PyResul #[pymethods(crate = "crate")] impl Coroutine { + #[getter] + fn __name__(&self) -> PyResult> { + self.name + .clone() + .ok_or_else(|| PyAttributeError::new_err("__name__")) + } + + #[getter] + fn __qualname__(&self) -> PyResult> { + self.qualname + .clone() + .ok_or_else(|| PyAttributeError::new_err("__qualname__")) + } + fn send(&mut self, py: Python<'_>, _value: &PyAny) -> PyResult { iter_result(self.poll(py, None)?) } @@ -135,3 +158,21 @@ impl Coroutine { self.poll(py, None) } } + +impl Drop for Coroutine { + fn drop(&mut self) { + if self.future.is_some() { + Python::with_gil(|gil| { + let qualname = self + .qualname + .as_ref() + .map_or(Ok(""), |n| n.as_ref(gil).to_str()) + .unwrap(); + let message = format!("coroutine {} was never awaited", qualname); + PyErr::warn(gil, gil.get_type::(), &message, 2) + .expect("warning error"); + self.poll(gil, None).expect("coroutine close error"); + }) + } + } +} diff --git a/src/impl_/coroutine.rs b/src/impl_/coroutine.rs index 843c42f169a..fca6625874d 100644 --- a/src/impl_/coroutine.rs +++ b/src/impl_/coroutine.rs @@ -1,19 +1,36 @@ -use crate::coroutine::Coroutine; -use crate::impl_::wrap::OkWrap; -use crate::{IntoPy, PyErr, PyObject, Python}; use std::future::Future; -/// Used to wrap the result of async `#[pyfunction]` and `#[pymethods]`. -pub fn wrap_future(future: F) -> Coroutine +use crate::{ + coroutine::Coroutine, + types::{PyModule, PyString}, + IntoPy, PyClass, PyErr, PyObject, Python, +}; + +pub fn new_coroutine(name: &PyString, qualname: &PyString, future: F) -> Coroutine where - F: Future + Send + 'static, - R: OkWrap, + F: Future> + Send + 'static, T: IntoPy, - PyErr: From, + PyErr: From, { - let future = async move { - // SAFETY: GIL is acquired when future is polled (see `Coroutine::poll`) - future.await.wrap(unsafe { Python::assume_gil_acquired() }) + Coroutine::new(Some(name.into()), Some(qualname.into()), future) +} + +pub fn coroutine_qualname<'py>( + py: Python<'py>, + module: Option<&PyModule>, + name: &str, +) -> &'py PyString { + match module.and_then(|m| m.name().ok()) { + Some(module) => PyString::new(py, &format!("{}.{}", module, name)), + None => PyString::new(py, name), + } +} + +pub fn method_coroutine_qualname<'py, T: PyClass>(py: Python<'py>, name: &str) -> &'py PyString { + let class = T::NAME; + let qualname = match T::MODULE { + Some(module) => format!("{}.{}.{}", module, class, name), + None => format!("{}.{}", class, name), }; - Coroutine::from_future(future) + PyString::new(py, &qualname) } diff --git a/src/impl_/wrap.rs b/src/impl_/wrap.rs index a73e3597302..dea7a0c287b 100644 --- a/src/impl_/wrap.rs +++ b/src/impl_/wrap.rs @@ -20,6 +20,7 @@ impl SomeWrap> for Option { /// Used to wrap the result of `#[pyfunction]` and `#[pymethods]`. pub trait OkWrap { type Error; + fn wrap_no_gil(self) -> Result; fn wrap(self, py: Python<'_>) -> Result, Self::Error>; } @@ -30,6 +31,9 @@ where T: IntoPy, { type Error = PyErr; + fn wrap_no_gil(self) -> Result { + Ok(self) + } fn wrap(self, py: Python<'_>) -> PyResult> { Ok(self.into_py(py)) } @@ -40,6 +44,9 @@ where T: IntoPy, { type Error = E; + fn wrap_no_gil(self) -> Result { + self + } fn wrap(self, py: Python<'_>) -> Result, Self::Error> { self.map(|o| o.into_py(py)) } @@ -57,4 +64,21 @@ mod tests { let b: Option = SomeWrap::wrap(None); assert_eq!(b, None); } + + #[test] + fn wrap_result() { + let a: PyResult = OkWrap::wrap_no_gil(42u8); + assert_eq!(a.unwrap(), 42); + + let b: Result = OkWrap::wrap_no_gil(Err("error")); + assert_eq!(b, Err("error")); + + Python::with_gil(|gil| { + let c: PyResult = OkWrap::wrap(gil.None(), gil); + assert!(c.unwrap().is_none(gil)); + + let d: Result = OkWrap::::wrap(Err("error"), gil); + assert_eq!(d.unwrap_err(), "error"); + }); + } } diff --git a/tests/test_coroutine.rs b/tests/test_coroutine.rs index 7c195e63733..c993ef69a14 100644 --- a/tests/test_coroutine.rs +++ b/tests/test_coroutine.rs @@ -1,8 +1,10 @@ #![cfg(feature = "macros")] #![cfg(not(target_arch = "wasm32"))] +use std::ops::Deref; use std::{task::Poll, thread, time::Duration}; use futures::{channel::oneshot, future::poll_fn}; +use pyo3::types::{IntoPyDict, PyType}; use pyo3::{prelude::*, py_run}; #[path = "../src/tests/common.rs"] @@ -30,6 +32,69 @@ fn noop_coroutine() { }) } +#[test] +fn test_coroutine_qualname() { + #[pyfunction] + async fn my_fn() {} + #[pyclass] + struct MyClass; + #[pymethods] + impl MyClass { + #[new] + fn new() -> Self { + Self + } + // TODO use &self when possible + async fn my_method(_self: Py) {} + #[classmethod] + async fn my_classmethod(_cls: Py) {} + #[staticmethod] + async fn my_staticmethod() {} + } + #[pyclass(module = "my_module")] + struct MyClassWithModule; + #[pymethods] + impl MyClassWithModule { + #[new] + fn new() -> Self { + Self + } + // TODO use &self when possible + async fn my_method(_self: Py) {} + #[classmethod] + async fn my_classmethod(_cls: Py) {} + #[staticmethod] + async fn my_staticmethod() {} + } + Python::with_gil(|gil| { + let test = r#" + for coro, name, qualname in [ + (my_fn(), "my_fn", "my_fn"), + (my_fn_with_module(), "my_fn", "my_module.my_fn"), + (MyClass().my_method(), "my_method", "MyClass.my_method"), + #(MyClass().my_classmethod(), "my_classmethod", "MyClass.my_classmethod"), + (MyClass.my_staticmethod(), "my_staticmethod", "MyClass.my_staticmethod"), + (MyClassWithModule().my_method(), "my_method", "my_module.MyClassWithModule.my_method"), + #(MyClassWithModule().my_classmethod(), "my_classmethod", "my_module.MyClassWithModule.my_classmethod"), + (MyClassWithModule.my_staticmethod(), "my_staticmethod", "my_module.MyClassWithModule.my_staticmethod"), + ]: + assert coro.__name__ == name and coro.__qualname__ == qualname + "#; + let my_module = PyModule::new(gil, "my_module").unwrap(); + let locals = [ + ("my_fn", wrap_pyfunction!(my_fn, gil).unwrap().deref()), + ( + "my_fn_with_module", + wrap_pyfunction!(my_fn, my_module).unwrap(), + ), + ("MyClass", gil.get_type::()), + ("MyClassWithModule", gil.get_type::()), + ] + .into_py_dict(gil); + py_run!(gil, *locals, &handle_windows(test)); + }) +} + #[test] fn sleep_0_like_coroutine() { #[pyfunction]